9importtorch.nnasnn10importtorch.utils.data11fromlabml\_helpers.moduleimportModule1213fromlabmlimportexperiment,tracker14fromlabml.configsimportoption15fromlabml\_helpers.datasets.mnistimportMNISTConfigs16fromlabml\_helpers.deviceimportDeviceConfigs17fromlabml\_helpers.metrics.accuracyimportAccuracy18fromlabml\_helpers.seedimportSeedConfigs19fromlabml\_helpers.train\_validimportTrainValidConfigs,BatchIndex,hook\_model\_outputs20fromlabml\_nn.optimizers.configsimportOptimizerConfigs

23classModel(Module):

27def\_\_init\_\_(self):28super().\_\_init\_\_()29self.conv1=nn.Conv2d(1,20,5,1)30self.pool1=nn.MaxPool2d(2)31self.conv2=nn.Conv2d(20,50,5,1)32self.pool2=nn.MaxPool2d(2)33self.fc1=nn.Linear(16\*50,500)34self.fc2=nn.Linear(500,10)35self.activation=nn.ReLU()

37defforward(self,x):38x=self.activation(self.conv1(x))39x=self.pool1(x)40x=self.activation(self.conv2(x))41x=self.pool2(x)42x=self.activation(self.fc1(x.view(-1,16\*50)))43returnself.fc2(x)

46classConfigs(MNISTConfigs,TrainValidConfigs):

50optimizer:torch.optim.Adam51model:nn.Module52set\_seed=SeedConfigs()53device:torch.device=DeviceConfigs()54epochs:int=105556is\_save\_models=True57model:nn.Module58inner\_iterations=105960accuracy\_func=Accuracy()61loss\_func=nn.CrossEntropyLoss()

63definit(self):64tracker.set\_queue("loss.\*",20,True)65tracker.set\_scalar("accuracy.\*",True)66hook\_model\_outputs(self.mode,self.model,'model')67self.state\_modules=[self.accuracy\_func]

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

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

74ifself.mode.is\_train:75tracker.add\_global\_step(len(data))

78withself.mode.update(is\_log\_activations=batch\_idx.is\_last):79output=self.model(data)

82loss=self.loss\_func(output,target)

84self.accuracy\_func(output,target)

86tracker.add("loss.",loss)

89ifself.mode.is\_train:

91loss.backward()

94self.optimizer.step()

96ifbatch\_idx.is\_last:97tracker.add('model',self.model)98tracker.add('optimizer',(self.optimizer,{'model':self.model}))

100self.optimizer.zero\_grad()

103tracker.save()

106@option(Configs.model)107defmodel(c:Configs):108returnModel().to(c.device)109110111@option(Configs.optimizer)112def\_optimizer(c:Configs):

117opt\_conf=OptimizerConfigs()118opt\_conf.parameters=c.model.parameters()119returnopt\_conf

122defmain():123conf=Configs()124conf.inner\_iterations=10125experiment.create(name='mnist\_ada\_belief')126experiment.configs(conf,{'inner\_iterations':10,

128'optimizer.optimizer':'Adam',129'optimizer.learning\_rate':1.5e-4})130conf.set\_seed.set()131experiment.add\_pytorch\_models(dict(model=conf.model))132withexperiment.start():133conf.run()134135136if\_\_name\_\_=='\_\_main\_\_':137main()