පොන්ඩර්නෙට් සමානාත්මතා කාර්ය අත්හදා බැලීම

13from typing import Any
14
15import torch
16from torch import nn
17from torch.utils.data import DataLoader
18
19from labml import tracker, experiment
20from labml_helpers.metrics.accuracy import AccuracyDirect
21from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
22from labml_nn.adaptive_computation.parity import ParityDataset
23from labml_nn.adaptive_computation.ponder_net import ParityPonderGRU, ReconstructionLoss, RegularizationLoss

26class Configs(SimpleTrainValidConfigs):

33    epochs: int = 100

35    n_batches: int = 500

37    batch_size: int = 128

40    model: ParityPonderGRU

43    loss_rec: ReconstructionLoss

45    loss_reg: RegularizationLoss

51    n_elems: int = 8

53    n_hidden: int = 64

55    max_steps: int = 20

58    lambda_p: float = 0.2

60    beta: float = 0.01

63    grad_norm_clip: float = 1.0

66    train_loader: DataLoader
67    valid_loader: DataLoader

70    accuracy = AccuracyDirect()

72    def init(self):

74        tracker.set_scalar('loss.*', True)
75        tracker.set_scalar('loss_reg.*', True)
76        tracker.set_scalar('accuracy.*', True)
77        tracker.set_scalar('steps.*', True)

80        self.state_modules = [self.accuracy]

83        self.model = ParityPonderGRU(self.n_elems, self.n_hidden, self.max_steps).to(self.device)

85        self.loss_rec = ReconstructionLoss(nn.BCEWithLogitsLoss(reduction='none')).to(self.device)

87        self.loss_reg = RegularizationLoss(self.lambda_p, self.max_steps).to(self.device)

90        self.train_loader = DataLoader(ParityDataset(self.batch_size * self.n_batches, self.n_elems),
91                                       batch_size=self.batch_size)
92        self.valid_loader = DataLoader(ParityDataset(self.batch_size * 32, self.n_elems),
93                                       batch_size=self.batch_size)

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

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

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

106        if self.mode.is_train:
107            tracker.add_global_step(len(data))

110        p, y_hat, p_sampled, y_hat_sampled = self.model(data)

113        loss_rec = self.loss_rec(p, y_hat, target.to(torch.float))
114        tracker.add("loss.", loss_rec)

117        loss_reg = self.loss_reg(p)
118        tracker.add("loss_reg.", loss_reg)

121        loss = loss_rec + self.beta * loss_reg

124        steps = torch.arange(1, p.shape[0] + 1, device=p.device)
125        expected_steps = (p * steps[:, None]).sum(dim=0)
126        tracker.add("steps.", expected_steps)

129        self.accuracy(y_hat_sampled > 0, target)
130
131        if self.mode.is_train:

133            loss.backward()

135            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)

137            self.optimizer.step()

139            self.optimizer.zero_grad()

141            tracker.save()

144def main():

148    experiment.create(name='ponder_net')
149
150    conf = Configs()
151    experiment.configs(conf, {
152        'optimizer.optimizer': 'Adam',
153        'optimizer.learning_rate': 0.0003,
154    })
155
156    with experiment.start():
157        conf.run()

160if __name__ == '__main__':
161    main()