15importtorch.nnasnn1617fromlabmlimportexperiment,logger18fromlabml.configsimportoption19fromlabml\_nn.experiments.cifar10importCIFAR10Configs,CIFAR10VGGModel20fromlabml\_nn.normalization.batch\_normimportBatchNorm

23classConfigs(CIFAR10Configs):

30pass

33classSmallModel(CIFAR10VGGModel):

40defconv\_block(self,in\_channels,out\_channels)-\>nn.Module:

44returnnn.Sequential(

46nn.Conv2d(in\_channels,out\_channels,kernel\_size=3,padding=1),

48BatchNorm(out\_channels,track\_running\_stats=False),

50nn.ReLU(inplace=True),51)

53def\_\_init\_\_(self):

55super().\_\_init\_\_([[32,32],[64,64],[128],[128],[128]])

58@option(Configs.model)59def\_small\_model(c:Configs):

63returnSmallModel().to(c.device)

66defmain():

68experiment.create(name='cifar10',comment='small model')

70conf=Configs()

72experiment.configs(conf,{73'optimizer.optimizer':'Adam',74'optimizer.learning\_rate':2.5e-4,75})

77experiment.add\_pytorch\_models({'model':conf.model})

79logger.inspect(params=(sum(p.numel()forpinconf.model.parameters()ifp.requires\_grad)))

81withexperiment.start():82conf.run()

86if\_\_name\_\_=='\_\_main\_\_':87main()