19importnumpyasnp20importtorch21importtorch.utils.data22importtorchvision.transforms.functional23fromtorchimportnn2425fromlabmlimportlab,tracker,experiment,monit26fromlabml.configsimportBaseConfigs27fromlabml\_helpers.deviceimportDeviceConfigs28fromlabml\_nn.unet.carvanaimportCarvanaDataset29fromlabml\_nn.unetimportUNet

32classConfigs(BaseConfigs):

39device:torch.device=DeviceConfigs()

42model:UNet

45image\_channels:int=3

47mask\_channels:int=1

50batch\_size:int=1

52learning\_rate:float=2.5e-4

55epochs:int=4

58dataset:CarvanaDataset

60data\_loader:torch.utils.data.DataLoader

63loss\_func=nn.BCELoss()

65sigmoid=nn.Sigmoid()

68optimizer:torch.optim.Adam

70definit(self):

72self.dataset=CarvanaDataset(lab.get\_data\_path()/'carvana'/'train',73lab.get\_data\_path()/'carvana'/'train\_masks')

75self.model=UNet(self.image\_channels,self.mask\_channels).to(self.device)

78self.data\_loader=torch.utils.data.DataLoader(self.dataset,self.batch\_size,79shuffle=True,pin\_memory=True)

81self.optimizer=torch.optim.Adam(self.model.parameters(),lr=self.learning\_rate)

84tracker.set\_image("sample",True)

[email protected]\_grad()87defsample(self,idx=-1):

93x,\_=self.dataset[np.random.randint(len(self.dataset))]

95x=x.to(self.device)

98mask=self.sigmoid(self.model(x[None,:]))

100x=torchvision.transforms.functional.center\_crop(x,[mask.shape[2],mask.shape[3]])

102tracker.save('sample',x\*mask)

104deftrain(self):

112for\_,(image,mask)inmonit.mix(('Train',self.data\_loader),(self.sample,list(range(50)))):

114tracker.add\_global\_step()

116image,mask=image.to(self.device),mask.to(self.device)

119self.optimizer.zero\_grad()

121logits=self.model(image)

124mask=torchvision.transforms.functional.center\_crop(mask,[logits.shape[2],logits.shape[3]])

126loss=self.loss\_func(self.sigmoid(logits),mask)

128loss.backward()

130self.optimizer.step()

132tracker.save('loss',loss)

134defrun(self):

138for\_inmonit.loop(self.epochs):

140self.train()

142tracker.new\_line()

144experiment.save\_checkpoint()

147defmain():

149experiment.create(name='unet')

152configs=Configs()

155experiment.configs(configs,{})

158configs.init()

161experiment.add\_pytorch\_models({'model':configs.model})

164withexperiment.start():165configs.run()

169if\_\_name\_\_=='\_\_main\_\_':170main()