14fromtypingimportAny1516importtorch.nnasnn17importtorch.utils.data1819fromlabmlimporttracker,experiment20fromlabml.configsimportoption,calculate21fromlabml\_helpers.moduleimportModule22fromlabml\_helpers.scheduleimportSchedule,RelativePiecewise23fromlabml\_helpers.train\_validimportBatchIndex24fromlabml\_nn.experiments.mnistimportMNISTConfigs25fromlabml\_nn.uncertainty.evidenceimportKLDivergenceLoss,TrackStatistics,MaximumLikelihoodLoss,\26CrossEntropyBayesRisk,SquaredErrorBayesRisk

29classModel(Module):

34def\_\_init\_\_(self,dropout:float):35super().\_\_init\_\_()

37self.conv1=nn.Conv2d(1,20,kernel\_size=5)

39self.act1=nn.ReLU()

41self.max\_pool1=nn.MaxPool2d(2,2)

43self.conv2=nn.Conv2d(20,50,kernel\_size=5)

45self.act2=nn.ReLU()

47self.max\_pool2=nn.MaxPool2d(2,2)

49self.fc1=nn.Linear(50\*4\*4,500)

51self.act3=nn.ReLU()

55self.fc2=nn.Linear(500,10)

57self.dropout=nn.Dropout(p=dropout)

59def\_\_call\_\_(self,x:torch.Tensor):

65x=self.max\_pool1(self.act1(self.conv1(x)))

68x=self.max\_pool2(self.act2(self.conv2(x)))

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

72x=self.act3(self.fc1(x))

74x=self.dropout(x)

76returnself.fc2(x)

79classConfigs(MNISTConfigs):

87kl\_div\_loss=KLDivergenceLoss()

89kl\_div\_coef:Schedule

91kl\_div\_coef\_schedule=[(0,0.),(0.2,0.01),(1,1.)]

93stats=TrackStatistics()

95dropout:float=0.5

97outputs\_to\_evidence:Module

99definit(self):

104tracker.set\_scalar("loss.\*",True)105tracker.set\_scalar("accuracy.\*",True)106tracker.set\_histogram('u.\*',True)107tracker.set\_histogram('prob.\*',False)108tracker.set\_scalar('annealing\_coef.\*',False)109tracker.set\_scalar('kl\_div\_loss.\*',False)

112self.state\_modules=[]

114defstep(self,batch:Any,batch\_idx:BatchIndex):

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

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

126eye=torch.eye(10).to(torch.float).to(self.device)127target=eye[target]

130ifself.mode.is\_train:131tracker.add\_global\_step(len(data))

134outputs=self.model(data)

136evidence=self.outputs\_to\_evidence(outputs)

139loss=self.loss\_func(evidence,target)

141kl\_div\_loss=self.kl\_div\_loss(evidence,target)142tracker.add("loss.",loss)143tracker.add("kl\_div\_loss.",kl\_div\_loss)

146annealing\_coef=min(1.,self.kl\_div\_coef(tracker.get\_global\_step()))147tracker.add("annealing\_coef.",annealing\_coef)

150loss=loss+annealing\_coef\*kl\_div\_loss

153self.stats(evidence,target)

156ifself.mode.is\_train:

158loss.backward()

160self.optimizer.step()

162self.optimizer.zero\_grad()

165tracker.save()

168@option(Configs.model)169defmnist\_model(c:Configs):

173returnModel(c.dropout).to(c.device)

176@option(Configs.kl\_div\_coef)177defkl\_div\_coef(c:Configs):

183returnRelativePiecewise(c.kl\_div\_coef\_schedule,c.epochs\*len(c.train\_dataset))

187calculate(Configs.loss\_func,'max\_likelihood\_loss',lambda:MaximumLikelihoodLoss())

189calculate(Configs.loss\_func,'cross\_entropy\_bayes\_risk',lambda:CrossEntropyBayesRisk())

191calculate(Configs.loss\_func,'squared\_error\_bayes\_risk',lambda:SquaredErrorBayesRisk())

194calculate(Configs.outputs\_to\_evidence,'relu',lambda:nn.ReLU())

196calculate(Configs.outputs\_to\_evidence,'softplus',lambda:nn.Softplus())

199defmain():

201experiment.create(name='evidence\_mnist')

203conf=Configs()

205experiment.configs(conf,{206'optimizer.optimizer':'Adam',207'optimizer.learning\_rate':0.001,208'optimizer.weight\_decay':0.005,

212'loss\_func':'squared\_error\_bayes\_risk',213214'outputs\_to\_evidence':'softplus',215216'dropout':0.5,217})

219withexperiment.start():220conf.run()

224if\_\_name\_\_=='\_\_main\_\_':225main()