12importtorch.nnasnn1314fromlabmlimportexperiment15fromlabml.configsimportoption16fromlabml\_helpers.moduleimportModule17fromlabml\_nn.experiments.cifar10importCIFAR10Configs18fromlabml\_nn.normalization.group\_normimportGroupNorm

21classModel(Module):

26def\_\_init\_\_(self,groups:int=32):27super().\_\_init\_\_()28layers=[]

30in\_channels=3

32forblockin[[64,64],[128,128],[256,256,256],[512,512,512],[512,512,512]]:

34forchannelsinblock:35layers+=[nn.Conv2d(in\_channels,channels,kernel\_size=3,padding=1),36GroupNorm(groups,channels),37nn.ReLU(inplace=True)]38in\_channels=channels

40layers+=[nn.MaxPool2d(kernel\_size=2,stride=2)]

43self.layers=nn.Sequential(\*layers)

45self.fc=nn.Linear(512,10)

47defforward(self,x):

49x=self.layers(x)

51x=x.view(x.shape[0],-1)

53returnself.fc(x)

56classConfigs(CIFAR10Configs):

58groups:int=16

61@option(Configs.model)62defmodel(c:Configs):

66returnModel(c.groups).to(c.device)

69defmain():

71experiment.create(name='cifar10',comment='group norm')

73conf=Configs()

75experiment.configs(conf,{76'optimizer.optimizer':'Adam',77'optimizer.learning\_rate':2.5e-4,78})

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

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