36importitertools37importrandom38importzipfile39fromtypingimportTuple4041importtorch42importtorch.nnasnn43importtorchvision.transformsastransforms44fromPILimportImage45fromtorch.utils.dataimportDataLoader,Dataset46fromtorchvision.transformsimportInterpolationMode47fromtorchvision.utilsimportmake\_grid4849fromlabmlimportlab,tracker,experiment,monit50fromlabml.configsimportBaseConfigs51fromlabml.utils.downloadimportdownload\_file52fromlabml.utils.pytorchimportget\_modules53fromlabml\_helpers.deviceimportDeviceConfigs54fromlabml\_helpers.moduleimportModule

57classGeneratorResNet(Module):

62def\_\_init\_\_(self,input\_channels:int,n\_residual\_blocks:int):63super().\_\_init\_\_()

71out\_features=6472layers=[73nn.Conv2d(input\_channels,out\_features,kernel\_size=7,padding=3,padding\_mode='reflect'),74nn.InstanceNorm2d(out\_features),75nn.ReLU(inplace=True),76]77in\_features=out\_features

81for\_inrange(2):82out\_features\*=283layers+=[84nn.Conv2d(in\_features,out\_features,kernel\_size=3,stride=2,padding=1),85nn.InstanceNorm2d(out\_features),86nn.ReLU(inplace=True),87]88in\_features=out\_features

92for\_inrange(n\_residual\_blocks):93layers+=[ResidualBlock(out\_features)]

97for\_inrange(2):98out\_features//=299layers+=[100nn.Upsample(scale\_factor=2),101nn.Conv2d(in\_features,out\_features,kernel\_size=3,stride=1,padding=1),102nn.InstanceNorm2d(out\_features),103nn.ReLU(inplace=True),104]105in\_features=out\_features

108layers+=[nn.Conv2d(out\_features,input\_channels,7,padding=3,padding\_mode='reflect'),nn.Tanh()]

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

114self.apply(weights\_init\_normal)

116defforward(self,x):117returnself.layers(x)

120classResidualBlock(Module):

125def\_\_init\_\_(self,in\_features:int):126super().\_\_init\_\_()127self.block=nn.Sequential(128nn.Conv2d(in\_features,in\_features,kernel\_size=3,padding=1,padding\_mode='reflect'),129nn.InstanceNorm2d(in\_features),130nn.ReLU(inplace=True),131nn.Conv2d(in\_features,in\_features,kernel\_size=3,padding=1,padding\_mode='reflect'),132nn.InstanceNorm2d(in\_features),133nn.ReLU(inplace=True),134)

136defforward(self,x:torch.Tensor):137returnx+self.block(x)

140classDiscriminator(Module):

145def\_\_init\_\_(self,input\_shape:Tuple[int,int,int]):146super().\_\_init\_\_()147channels,height,width=input\_shape

151self.output\_shape=(1,height//2\*\*4,width//2\*\*4)152153self.layers=nn.Sequential(

155DiscriminatorBlock(channels,64,normalize=False),156DiscriminatorBlock(64,128),157DiscriminatorBlock(128,256),158DiscriminatorBlock(256,512),

161nn.ZeroPad2d((1,0,1,0)),162nn.Conv2d(512,1,kernel\_size=4,padding=1)163)

166self.apply(weights\_init\_normal)

168defforward(self,img):169returnself.layers(img)

172classDiscriminatorBlock(Module):

180def\_\_init\_\_(self,in\_filters:int,out\_filters:int,normalize:bool=True):181super().\_\_init\_\_()182layers=[nn.Conv2d(in\_filters,out\_filters,kernel\_size=4,stride=2,padding=1)]183ifnormalize:184layers.append(nn.InstanceNorm2d(out\_filters))185layers.append(nn.LeakyReLU(0.2,inplace=True))186self.layers=nn.Sequential(\*layers)

188defforward(self,x:torch.Tensor):189returnself.layers(x)

192defweights\_init\_normal(m):

196classname=m.\_\_class\_\_.\_\_name\_\_197ifclassname.find("Conv")!=-1:198torch.nn.init.normal\_(m.weight.data,0.0,0.02)

201defload\_image(path:str):

205image=Image.open(path)206ifimage.mode!='RGB':207image=Image.new("RGB",image.size).paste(image)208209returnimage

212classImageDataset(Dataset):

217@staticmethod218defdownload(dataset\_name:str):

223url=f'https://people.eecs.berkeley.edu/~taesung\_park/CycleGAN/datasets/{dataset\_name}.zip'

225root=lab.get\_data\_path()/'cycle\_gan'226ifnotroot.exists():227root.mkdir(parents=True)

229archive=root/f'{dataset\_name}.zip'

231download\_file(url,archive)

233withzipfile.ZipFile(archive,'r')asf:234f.extractall(root)

236def\_\_init\_\_(self,dataset\_name:str,transforms\_,mode:str):

245root=lab.get\_data\_path()/'cycle\_gan'/dataset\_name

247ifnotroot.exists():248self.download(dataset\_name)

251self.transform=transforms.Compose(transforms\_)

254path\_a=root/f'{mode}A'255path\_b=root/f'{mode}B'256self.files\_a=sorted(str(f)forfinpath\_a.iterdir())257self.files\_b=sorted(str(f)forfinpath\_b.iterdir())

259def\_\_getitem\_\_(self,index):

263return{"x":self.transform(load\_image(self.files\_a[index%len(self.files\_a)])),264"y":self.transform(load\_image(self.files\_b[index%len(self.files\_b)]))}

266def\_\_len\_\_(self):

268returnmax(len(self.files\_a),len(self.files\_b))

271classReplayBuffer:

285def\_\_init\_\_(self,max\_size:int=50):286self.max\_size=max\_size287self.data=[]

289defpush\_and\_pop(self,data:torch.Tensor):

291data=data.detach()292res=[]293forelementindata:294iflen(self.data)\<self.max\_size:295self.data.append(element)296res.append(element)297else:298ifrandom.uniform(0,1)\>0.5:299i=random.randint(0,self.max\_size-1)300res.append(self.data[i].clone())301self.data[i]=element302else:303res.append(element)304returntorch.stack(res)

307classConfigs(BaseConfigs):

311device:torch.device=DeviceConfigs()

314epochs:int=200315dataset\_name:str='monet2photo'316batch\_size:int=1317318data\_loader\_workers=8319320learning\_rate=0.0002321adam\_betas=(0.5,0.999)322decay\_start=100

326gan\_loss=torch.nn.MSELoss()

329cycle\_loss=torch.nn.L1Loss()330identity\_loss=torch.nn.L1Loss()

333img\_height=256334img\_width=256335img\_channels=3

338n\_residual\_blocks=9

341cyclic\_loss\_coefficient=10.0342identity\_loss\_coefficient=5.343344sample\_interval=500

347generator\_xy:GeneratorResNet348generator\_yx:GeneratorResNet349discriminator\_x:Discriminator350discriminator\_y:Discriminator

353generator\_optimizer:torch.optim.Adam354discriminator\_optimizer:torch.optim.Adam

357generator\_lr\_scheduler:torch.optim.lr\_scheduler.LambdaLR358discriminator\_lr\_scheduler:torch.optim.lr\_scheduler.LambdaLR

361dataloader:DataLoader362valid\_dataloader:DataLoader

364defsample\_images(self,n:int):

366batch=next(iter(self.valid\_dataloader))367self.generator\_xy.eval()368self.generator\_yx.eval()369withtorch.no\_grad():370data\_x,data\_y=batch['x'].to(self.generator\_xy.device),batch['y'].to(self.generator\_yx.device)371gen\_y=self.generator\_xy(data\_x)372gen\_x=self.generator\_yx(data\_y)

375data\_x=make\_grid(data\_x,nrow=5,normalize=True)376data\_y=make\_grid(data\_y,nrow=5,normalize=True)377gen\_x=make\_grid(gen\_x,nrow=5,normalize=True)378gen\_y=make\_grid(gen\_y,nrow=5,normalize=True)

381image\_grid=torch.cat((data\_x,gen\_y,data\_y,gen\_x),1)

384plot\_image(image\_grid)

386definitialize(self):

390input\_shape=(self.img\_channels,self.img\_height,self.img\_width)

393self.generator\_xy=GeneratorResNet(self.img\_channels,self.n\_residual\_blocks).to(self.device)394self.generator\_yx=GeneratorResNet(self.img\_channels,self.n\_residual\_blocks).to(self.device)395self.discriminator\_x=Discriminator(input\_shape).to(self.device)396self.discriminator\_y=Discriminator(input\_shape).to(self.device)

399self.generator\_optimizer=torch.optim.Adam(400itertools.chain(self.generator\_xy.parameters(),self.generator\_yx.parameters()),401lr=self.learning\_rate,betas=self.adam\_betas)402self.discriminator\_optimizer=torch.optim.Adam(403itertools.chain(self.discriminator\_x.parameters(),self.discriminator\_y.parameters()),404lr=self.learning\_rate,betas=self.adam\_betas)

409decay\_epochs=self.epochs-self.decay\_start410self.generator\_lr\_scheduler=torch.optim.lr\_scheduler.LambdaLR(411self.generator\_optimizer,lr\_lambda=lambdae:1.0-max(0,e-self.decay\_start)/decay\_epochs)412self.discriminator\_lr\_scheduler=torch.optim.lr\_scheduler.LambdaLR(413self.discriminator\_optimizer,lr\_lambda=lambdae:1.0-max(0,e-self.decay\_start)/decay\_epochs)

416transforms\_=[417transforms.Resize(int(self.img\_height\*1.12),InterpolationMode.BICUBIC),418transforms.RandomCrop((self.img\_height,self.img\_width)),419transforms.RandomHorizontalFlip(),420transforms.ToTensor(),421transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)),422]

425self.dataloader=DataLoader(426ImageDataset(self.dataset\_name,transforms\_,'train'),427batch\_size=self.batch\_size,428shuffle=True,429num\_workers=self.data\_loader\_workers,430)

433self.valid\_dataloader=DataLoader(434ImageDataset(self.dataset\_name,transforms\_,"test"),435batch\_size=5,436shuffle=True,437num\_workers=self.data\_loader\_workers,438)

440defrun(self):

542gen\_x\_buffer=ReplayBuffer()543gen\_y\_buffer=ReplayBuffer()

546forepochinmonit.loop(self.epochs):

548fori,batchinmonit.enum('Train',self.dataloader):

550data\_x,data\_y=batch['x'].to(self.device),batch['y'].to(self.device)

553true\_labels=torch.ones(data\_x.size(0),\*self.discriminator\_x.output\_shape,554device=self.device,requires\_grad=False)

556false\_labels=torch.zeros(data\_x.size(0),\*self.discriminator\_x.output\_shape,557device=self.device,requires\_grad=False)

561gen\_x,gen\_y=self.optimize\_generators(data\_x,data\_y,true\_labels)

564self.optimize\_discriminator(data\_x,data\_y,565gen\_x\_buffer.push\_and\_pop(gen\_x),gen\_y\_buffer.push\_and\_pop(gen\_y),566true\_labels,false\_labels)

569tracker.save()570tracker.add\_global\_step(max(len(data\_x),len(data\_y)))

573batches\_done=epoch\*len(self.dataloader)+i574ifbatches\_done%self.sample\_interval==0:

576experiment.save\_checkpoint()

578self.sample\_images(batches\_done)

581self.generator\_lr\_scheduler.step()582self.discriminator\_lr\_scheduler.step()

584tracker.new\_line()

586defoptimize\_generators(self,data\_x:torch.Tensor,data\_y:torch.Tensor,true\_labels:torch.Tensor):

592self.generator\_xy.train()593self.generator\_yx.train()

598loss\_identity=(self.identity\_loss(self.generator\_yx(data\_x),data\_x)+599self.identity\_loss(self.generator\_xy(data\_y),data\_y))

602gen\_y=self.generator\_xy(data\_x)603gen\_x=self.generator\_yx(data\_y)

608loss\_gan=(self.gan\_loss(self.discriminator\_y(gen\_y),true\_labels)+609self.gan\_loss(self.discriminator\_x(gen\_x),true\_labels))

616loss\_cycle=(self.cycle\_loss(self.generator\_yx(gen\_y),data\_x)+617self.cycle\_loss(self.generator\_xy(gen\_x),data\_y))

620loss\_generator=(loss\_gan+621self.cyclic\_loss\_coefficient\*loss\_cycle+622self.identity\_loss\_coefficient\*loss\_identity)

625self.generator\_optimizer.zero\_grad()626loss\_generator.backward()627self.generator\_optimizer.step()

630tracker.add({'loss.generator':loss\_generator,631'loss.generator.cycle':loss\_cycle,632'loss.generator.gan':loss\_gan,633'loss.generator.identity':loss\_identity})

636returngen\_x,gen\_y

638defoptimize\_discriminator(self,data\_x:torch.Tensor,data\_y:torch.Tensor,639gen\_x:torch.Tensor,gen\_y:torch.Tensor,640true\_labels:torch.Tensor,false\_labels:torch.Tensor):

653loss\_discriminator=(self.gan\_loss(self.discriminator\_x(data\_x),true\_labels)+654self.gan\_loss(self.discriminator\_x(gen\_x),false\_labels)+655self.gan\_loss(self.discriminator\_y(data\_y),true\_labels)+656self.gan\_loss(self.discriminator\_y(gen\_y),false\_labels))

659self.discriminator\_optimizer.zero\_grad()660loss\_discriminator.backward()661self.discriminator\_optimizer.step()

664tracker.add({'loss.discriminator':loss\_discriminator})

667deftrain():

672conf=Configs()

674experiment.create(name='cycle\_gan')

677experiment.configs(conf,{'dataset\_name':'summer2winter\_yosemite'})678conf.initialize()

683experiment.add\_pytorch\_models(get\_modules(conf))

685withexperiment.start():

687conf.run()

690defplot\_image(img:torch.Tensor):

694frommatplotlibimportpyplotasplt

697img=img.cpu()

699img\_min,img\_max=img.min(),img.max()

701img=(img-img\_min)/(img\_max-img\_min+1e-5)

703img=img.permute(1,2,0)

705plt.imshow(img)

707plt.axis('off')

709plt.show()

712defevaluate():

717trained\_run\_uuid='f73c1164184711eb9190b74249275441'

719conf=Configs()

721experiment.create(name='cycle\_gan\_inference')

723conf\_dict=experiment.load\_configs(trained\_run\_uuid)

732experiment.configs(conf,conf\_dict)733conf.initialize()

738experiment.add\_pytorch\_models(get\_modules(conf))

741experiment.load(trained\_run\_uuid)

744withexperiment.start():

746transforms\_=[747transforms.ToTensor(),748transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)),749]

755dataset=ImageDataset(conf.dataset\_name,transforms\_,'train')

757x\_image=dataset[10]['x']

759plot\_image(x\_image)

762conf.generator\_xy.eval()763conf.generator\_yx.eval()

766withtorch.no\_grad():

768data=x\_image.unsqueeze(0).to(conf.device)769generated\_y=conf.generator\_xy(data)

772plot\_image(generated\_y[0].cpu())773774775if\_\_name\_\_=='\_\_main\_\_':776train()