කැප්සියුලජාල සමඟ MNIST ඉලක්කම් වර්ගීකරණය කරන්න

14from typing import Any
15
16import torch.nn as nn
17import torch.nn.functional as F
18import torch.utils.data
19
20from labml import experiment, tracker
21from labml.configs import option
22from labml_helpers.datasets.mnist import MNISTConfigs
23from labml_helpers.metrics.accuracy import AccuracyDirect
24from labml_helpers.module import Module
25from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
26from labml_nn.capsule_networks import Squash, Router, MarginLoss

29class MNISTCapsuleNetworkModel(Module):

34    def __init__(self):
35        super().__init__()

37        self.conv1 = nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9, stride=1)

43        self.conv2 = nn.Conv2d(in_channels=256, out_channels=32 * 8, kernel_size=9, stride=2, padding=0)
44        self.squash = Squash()

50        self.digit_capsules = Router(32 * 6 * 6, 10, 8, 16, 3)

55        self.decoder = nn.Sequential(
56            nn.Linear(16 * 10, 512),
57            nn.ReLU(),
58            nn.Linear(512, 1024),
59            nn.ReLU(),
60            nn.Linear(1024, 784),
61            nn.Sigmoid()
62        )

64    def forward(self, data: torch.Tensor):

70        x = F.relu(self.conv1(data))

74        x = self.conv2(x)

77        caps = x.view(x.shape[0], 8, 32 * 6 * 6).permute(0, 2, 1)

79        caps = self.squash(caps)

82        caps = self.digit_capsules(caps)

85        with torch.no_grad():

87            pred = (caps ** 2).sum(-1).argmax(-1)

89            mask = torch.eye(10, device=data.device)[pred]

93        reconstructions = self.decoder((caps * mask[:, :, None]).view(x.shape[0], -1))

95        reconstructions = reconstructions.view(-1, 1, 28, 28)
96
97        return caps, reconstructions, pred

100class Configs(MNISTConfigs, SimpleTrainValidConfigs):

104    epochs: int = 10
105    model: nn.Module = 'capsule_network_model'
106    reconstruction_loss = nn.MSELoss()
107    margin_loss = MarginLoss(n_labels=10)
108    accuracy = AccuracyDirect()

110    def init(self):

112        tracker.set_scalar('loss.*', True)
113        tracker.set_scalar('accuracy.*', True)

116        self.state_modules = [self.accuracy]

118    def step(self, batch: Any, batch_idx: BatchIndex):

123        self.model.train(self.mode.is_train)

126        data, target = batch[0].to(self.device), batch[1].to(self.device)

129        if self.mode.is_train:
130            tracker.add_global_step(len(data))

133        with self.mode.update(is_log_activations=batch_idx.is_last):

135            caps, reconstructions, pred = self.model(data)

138        loss = self.margin_loss(caps, target) + 0.0005 * self.reconstruction_loss(reconstructions, data)
139        tracker.add("loss.", loss)

142        self.accuracy(pred, target)
143
144        if self.mode.is_train:
145            loss.backward()
146
147            self.optimizer.step()

149            if batch_idx.is_last:
150                tracker.add('model', self.model)
151            self.optimizer.zero_grad()
152
153            tracker.save()

156@option(Configs.model)
157def capsule_network_model(c: Configs):

159    return MNISTCapsuleNetworkModel().to(c.device)

162def main():

166    experiment.create(name='capsule_network_mnist')
167    conf = Configs()
168    experiment.configs(conf, {'optimizer.optimizer': 'Adam',
169                              'optimizer.learning_rate': 1e-3})
170
171    experiment.add_pytorch_models({'model': conf.model})
172
173    with experiment.start():
174        conf.run()
175
176
177if __name__ == '__main__':
178    main()