11importtorch.nnasnn12importtorch.utils.data13fromlabml\_helpers.moduleimportModule1415fromlabmlimporttracker16fromlabml.configsimportoption17fromlabml\_helpers.datasets.mnistimportMNISTConfigsasMNISTDatasetConfigs18fromlabml\_helpers.deviceimportDeviceConfigs19fromlabml\_helpers.metrics.accuracyimportAccuracy20fromlabml\_helpers.train\_validimportTrainValidConfigs,BatchIndex,hook\_model\_outputs21fromlabml\_nn.optimizers.configsimportOptimizerConfigs

24classMNISTConfigs(MNISTDatasetConfigs,TrainValidConfigs):

32optimizer:torch.optim.Adam

34device:torch.device=DeviceConfigs()

37model:Module

39epochs:int=10

42inner\_iterations=10

45accuracy=Accuracy()

47loss\_func=nn.CrossEntropyLoss()

49definit(self):

54tracker.set\_scalar("loss.\*",True)55tracker.set\_scalar("accuracy.\*",True)

57hook\_model\_outputs(self.mode,self.model,'model')

62self.state\_modules=[self.accuracy]

64defstep(self,batch:any,batch\_idx:BatchIndex):

70self.model.train(self.mode.is\_train)

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

76ifself.mode.is\_train:77tracker.add\_global\_step(len(data))

80withself.mode.update(is\_log\_activations=batch\_idx.is\_last):

82output=self.model(data)

85loss=self.loss\_func(output,target)86tracker.add("loss.",loss)

89self.accuracy(output,target)90self.accuracy.track()

93ifself.mode.is\_train:

95loss.backward()

97self.optimizer.step()

99ifbatch\_idx.is\_last:100tracker.add('model',self.model)

102self.optimizer.zero\_grad()

105tracker.save()

108@option(MNISTConfigs.optimizer)109def\_optimizer(c:MNISTConfigs):

113opt\_conf=OptimizerConfigs()114opt\_conf.parameters=c.model.parameters()115opt\_conf.optimizer='Adam'116returnopt\_conf