13fromtypingimportAny1415importtorch16fromtorchimportnn17fromtorch.utils.dataimportDataLoader1819fromlabmlimporttracker,experiment20fromlabml\_helpers.metrics.accuracyimportAccuracyDirect21fromlabml\_helpers.train\_validimportSimpleTrainValidConfigs,BatchIndex22fromlabml\_nn.adaptive\_computation.parityimportParityDataset23fromlabml\_nn.adaptive\_computation.ponder\_netimportParityPonderGRU,ReconstructionLoss,RegularizationLoss

26classConfigs(SimpleTrainValidConfigs):

33epochs:int=100

35n\_batches:int=500

37batch\_size:int=128

40model:ParityPonderGRU

43loss\_rec:ReconstructionLoss

45loss\_reg:RegularizationLoss

51n\_elems:int=8

53n\_hidden:int=64

55max\_steps:int=20

58lambda\_p:float=0.2

60beta:float=0.01

63grad\_norm\_clip:float=1.0

66train\_loader:DataLoader67valid\_loader:DataLoader

70accuracy=AccuracyDirect()

72definit(self):

74tracker.set\_scalar('loss.\*',True)75tracker.set\_scalar('loss\_reg.\*',True)76tracker.set\_scalar('accuracy.\*',True)77tracker.set\_scalar('steps.\*',True)

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

83self.model=ParityPonderGRU(self.n\_elems,self.n\_hidden,self.max\_steps).to(self.device)

85self.loss\_rec=ReconstructionLoss(nn.BCEWithLogitsLoss(reduction='none')).to(self.device)

87self.loss\_reg=RegularizationLoss(self.lambda\_p,self.max\_steps).to(self.device)

90self.train\_loader=DataLoader(ParityDataset(self.batch\_size\*self.n\_batches,self.n\_elems),91batch\_size=self.batch\_size)92self.valid\_loader=DataLoader(ParityDataset(self.batch\_size\*32,self.n\_elems),93batch\_size=self.batch\_size)

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

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

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

106ifself.mode.is\_train:107tracker.add\_global\_step(len(data))

110p,y\_hat,p\_sampled,y\_hat\_sampled=self.model(data)

113loss\_rec=self.loss\_rec(p,y\_hat,target.to(torch.float))114tracker.add("loss.",loss\_rec)

117loss\_reg=self.loss\_reg(p)118tracker.add("loss\_reg.",loss\_reg)

121loss=loss\_rec+self.beta\*loss\_reg

124steps=torch.arange(1,p.shape[0]+1,device=p.device)125expected\_steps=(p\*steps[:,None]).sum(dim=0)126tracker.add("steps.",expected\_steps)

129self.accuracy(y\_hat\_sampled\>0,target)130131ifself.mode.is\_train:

133loss.backward()

135torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=self.grad\_norm\_clip)

137self.optimizer.step()

139self.optimizer.zero\_grad()

141tracker.save()

144defmain():

148experiment.create(name='ponder\_net')149150conf=Configs()151experiment.configs(conf,{152'optimizer.optimizer':'Adam',153'optimizer.learning\_rate':0.0003,154})155156withexperiment.start():157conf.run()

160if\_\_name\_\_=='\_\_main\_\_':161main()