ස්නායුකජාලයක දැනුම ආසවනය කිරීම

74import torch
75import torch.nn.functional
76from torch import nn
77
78from labml import experiment, tracker
79from labml.configs import option
80from labml_helpers.train_valid import BatchIndex
81from labml_nn.distillation.large import LargeModel
82from labml_nn.distillation.small import SmallModel
83from labml_nn.experiments.cifar10 import CIFAR10Configs

86class Configs(CIFAR10Configs):

94    model: SmallModel

96    large: LargeModel

98    kl_div_loss = nn.KLDivLoss(log_target=True)

100    loss_func = nn.CrossEntropyLoss()

102    temperature: float = 5.

108    soft_targets_weight: float = 100.

110    label_loss_weight: float = 0.5

112    def step(self, batch: any, batch_idx: BatchIndex):

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

122        self.large.eval()

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

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

132        with torch.no_grad():
133            large_logits = self.large(data)

136        output = self.model(data)

140        soft_targets = nn.functional.log_softmax(large_logits / self.temperature, dim=-1)

143        soft_prob = nn.functional.log_softmax(output / self.temperature, dim=-1)

146        soft_targets_loss = self.kl_div_loss(soft_prob, soft_targets)

148        label_loss = self.loss_func(output, target)

150        loss = self.soft_targets_weight * soft_targets_loss + self.label_loss_weight * label_loss

152        tracker.add({"loss.kl_div.": soft_targets_loss,
153                     "loss.nll": label_loss,
154                     "loss.": loss})

157        self.accuracy(output, target)
158        self.accuracy.track()

161        if self.mode.is_train:

163            loss.backward()

165            self.optimizer.step()

167            if batch_idx.is_last:
168                tracker.add('model', self.model)

170            self.optimizer.zero_grad()

173        tracker.save()

176@option(Configs.large)
177def _large_model(c: Configs):

181    return LargeModel().to(c.device)

184@option(Configs.model)
185def _small_student_model(c: Configs):

189    return SmallModel().to(c.device)

192def get_saved_model(run_uuid: str, checkpoint: int):

197    from labml_nn.distillation.large import Configs as LargeConfigs

200    experiment.evaluate()

202    conf = LargeConfigs()

204    experiment.configs(conf, experiment.load_configs(run_uuid))

206    experiment.add_pytorch_models({'model': conf.model})

208    experiment.load(run_uuid, checkpoint)

210    experiment.start()

213    return conf.model

216def main(run_uuid: str, checkpoint: int):

221    large_model = get_saved_model(run_uuid, checkpoint)

223    experiment.create(name='distillation', comment='cifar10')

225    conf = Configs()

227    conf.large = large_model

229    experiment.configs(conf, {
230        'optimizer.optimizer': 'Adam',
231        'optimizer.learning_rate': 2.5e-4,
232        'model': '_small_student_model',
233    })

235    experiment.add_pytorch_models({'model': conf.model})

237    experiment.load(None, None)

239    with experiment.start():
240        conf.run()

244if __name__ == '__main__':
245    main('d46cd53edaec11eb93c38d6538aee7d6', 1_000_000)