31importmath32frompathlibimportPath33fromtypingimportIterator,Tuple3435importtorch36importtorch.utils.data37importtorchvision38fromPILimportImage3940fromlabmlimporttracker,lab,monit,experiment41fromlabml.configsimportBaseConfigs42fromlabml\_helpers.deviceimportDeviceConfigs43fromlabml\_helpers.train\_validimportModeState,hook\_model\_outputs44fromlabml\_nn.gan.styleganimportDiscriminator,Generator,MappingNetwork,GradientPenalty,PathLengthPenalty45fromlabml\_nn.gan.wassersteinimportDiscriminatorLoss,GeneratorLoss46fromlabml\_nn.utilsimportcycle\_dataloader

49classDataset(torch.utils.data.Dataset):

56def\_\_init\_\_(self,path:str,image\_size:int):

61super().\_\_init\_\_()

64self.paths=[pforpinPath(path).glob(f'\*\*/\*.jpg')]

67self.transform=torchvision.transforms.Compose([

69torchvision.transforms.Resize(image\_size),

71torchvision.transforms.ToTensor(),72])

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

76returnlen(self.paths)

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

80path=self.paths[index]81img=Image.open(path)82returnself.transform(img)

85classConfigs(BaseConfigs):

93device:torch.device=DeviceConfigs()

96discriminator:Discriminator

98generator:Generator

100mapping\_network:MappingNetwork

104discriminator\_loss:DiscriminatorLoss105generator\_loss:GeneratorLoss

108generator\_optimizer:torch.optim.Adam109discriminator\_optimizer:torch.optim.Adam110mapping\_network\_optimizer:torch.optim.Adam

113gradient\_penalty=GradientPenalty()

115gradient\_penalty\_coefficient:float=10.

118path\_length\_penalty:PathLengthPenalty

121loader:Iterator

124batch\_size:int=32

126d\_latent:int=512

128image\_size:int=32

130mapping\_network\_layers:int=8

132learning\_rate:float=1e-3

134mapping\_network\_learning\_rate:float=1e-5

136gradient\_accumulate\_steps:int=1

138adam\_betas:Tuple[float,float]=(0.0,0.99)

140style\_mixing\_prob:float=0.9

143training\_steps:int=150\_000

146n\_gen\_blocks:int

154lazy\_gradient\_penalty\_interval:int=4

156lazy\_path\_penalty\_interval:int=32

158lazy\_path\_penalty\_after:int=5\_000

161log\_generated\_interval:int=500

163save\_checkpoint\_interval:int=2\_000

166mode:ModeState

168log\_layer\_outputs:bool=False

175dataset\_path:str=str(lab.get\_data\_path()/'stylegan2')

177definit(self):

182dataset=Dataset(self.dataset\_path,self.image\_size)

184dataloader=torch.utils.data.DataLoader(dataset,batch\_size=self.batch\_size,num\_workers=8,185shuffle=True,drop\_last=True,pin\_memory=True)

187self.loader=cycle\_dataloader(dataloader)

190log\_resolution=int(math.log2(self.image\_size))

193self.discriminator=Discriminator(log\_resolution).to(self.device)194self.generator=Generator(log\_resolution,self.d\_latent).to(self.device)

196self.n\_gen\_blocks=self.generator.n\_blocks

198self.mapping\_network=MappingNetwork(self.d\_latent,self.mapping\_network\_layers).to(self.device)

200self.path\_length\_penalty=PathLengthPenalty(0.99).to(self.device)

203ifself.log\_layer\_outputs:204hook\_model\_outputs(self.mode,self.discriminator,'discriminator')205hook\_model\_outputs(self.mode,self.generator,'generator')206hook\_model\_outputs(self.mode,self.mapping\_network,'mapping\_network')

209self.discriminator\_loss=DiscriminatorLoss().to(self.device)210self.generator\_loss=GeneratorLoss().to(self.device)

213self.discriminator\_optimizer=torch.optim.Adam(214self.discriminator.parameters(),215lr=self.learning\_rate,betas=self.adam\_betas216)217self.generator\_optimizer=torch.optim.Adam(218self.generator.parameters(),219lr=self.learning\_rate,betas=self.adam\_betas220)221self.mapping\_network\_optimizer=torch.optim.Adam(222self.mapping\_network.parameters(),223lr=self.mapping\_network\_learning\_rate,betas=self.adam\_betas224)

227tracker.set\_image("generated",True)

229defget\_w(self,batch\_size:int):

243iftorch.rand(()).item()\<self.style\_mixing\_prob:

245cross\_over\_point=int(torch.rand(()).item()\*self.n\_gen\_blocks)

247z2=torch.randn(batch\_size,self.d\_latent).to(self.device)248z1=torch.randn(batch\_size,self.d\_latent).to(self.device)

250w1=self.mapping\_network(z1)251w2=self.mapping\_network(z2)

253w1=w1[None,:,:].expand(cross\_over\_point,-1,-1)254w2=w2[None,:,:].expand(self.n\_gen\_blocks-cross\_over\_point,-1,-1)255returntorch.cat((w1,w2),dim=0)

257else:

259z=torch.randn(batch\_size,self.d\_latent).to(self.device)

261w=self.mapping\_network(z)

263returnw[None,:,:].expand(self.n\_gen\_blocks,-1,-1)

265defget\_noise(self,batch\_size:int):

272noise=[]

274resolution=4

277foriinrange(self.n\_gen\_blocks):

279ifi==0:280n1=None

282else:283n1=torch.randn(batch\_size,1,resolution,resolution,device=self.device)

285n2=torch.randn(batch\_size,1,resolution,resolution,device=self.device)

288noise.append((n1,n2))

291resolution\*=2

294returnnoise

296defgenerate\_images(self,batch\_size:int):

304w=self.get\_w(batch\_size)

306noise=self.get\_noise(batch\_size)

309images=self.generator(w,noise)

312returnimages,w

314defstep(self,idx:int):

320withmonit.section('Discriminator'):

322self.discriminator\_optimizer.zero\_grad()

325foriinrange(self.gradient\_accumulate\_steps):

327withself.mode.update(is\_log\_activations=(idx+1)%self.log\_generated\_interval==0):

329generated\_images,\_=self.generate\_images(self.batch\_size)

331fake\_output=self.discriminator(generated\_images.detach())

334real\_images=next(self.loader).to(self.device)

336if(idx+1)%self.lazy\_gradient\_penalty\_interval==0:337real\_images.requires\_grad\_()

339real\_output=self.discriminator(real\_images)

342real\_loss,fake\_loss=self.discriminator\_loss(real\_output,fake\_output)343disc\_loss=real\_loss+fake\_loss

346if(idx+1)%self.lazy\_gradient\_penalty\_interval==0:

348gp=self.gradient\_penalty(real\_images,real\_output)349tracker.add('loss.gp',gp)

351disc\_loss=disc\_loss+0.5\*self.gradient\_penalty\_coefficient\*gp\*self.lazy\_gradient\_penalty\_interval

354disc\_loss.backward()

357tracker.add('loss.discriminator',disc\_loss)358359if(idx+1)%self.log\_generated\_interval==0:

361tracker.add('discriminator',self.discriminator)

364torch.nn.utils.clip\_grad\_norm\_(self.discriminator.parameters(),max\_norm=1.0)

366self.discriminator\_optimizer.step()

369withmonit.section('Generator'):

371self.generator\_optimizer.zero\_grad()372self.mapping\_network\_optimizer.zero\_grad()

375foriinrange(self.gradient\_accumulate\_steps):

377generated\_images,w=self.generate\_images(self.batch\_size)

379fake\_output=self.discriminator(generated\_images)

382gen\_loss=self.generator\_loss(fake\_output)

385ifidx\>self.lazy\_path\_penalty\_afterand(idx+1)%self.lazy\_path\_penalty\_interval==0:

387plp=self.path\_length\_penalty(w,generated\_images)

389ifnottorch.isnan(plp):390tracker.add('loss.plp',plp)391gen\_loss=gen\_loss+plp

394gen\_loss.backward()

397tracker.add('loss.generator',gen\_loss)398399if(idx+1)%self.log\_generated\_interval==0:

401tracker.add('generator',self.generator)402tracker.add('mapping\_network',self.mapping\_network)

405torch.nn.utils.clip\_grad\_norm\_(self.generator.parameters(),max\_norm=1.0)406torch.nn.utils.clip\_grad\_norm\_(self.mapping\_network.parameters(),max\_norm=1.0)

409self.generator\_optimizer.step()410self.mapping\_network\_optimizer.step()

413if(idx+1)%self.log\_generated\_interval==0:414tracker.add('generated',torch.cat([generated\_images[:6],real\_images[:3]],dim=0))

416if(idx+1)%self.save\_checkpoint\_interval==0:417experiment.save\_checkpoint()

420tracker.save()

422deftrain(self):

428foriinmonit.loop(self.training\_steps):

430self.step(i)

432if(i+1)%self.log\_generated\_interval==0:433tracker.new\_line()

436defmain():

442experiment.create(name='stylegan2')

444configs=Configs()

447experiment.configs(configs,{448'device.cuda\_device':0,449'image\_size':64,450'log\_generated\_interval':200451})

454configs.init()

456experiment.add\_pytorch\_models(mapping\_network=configs.mapping\_network,457generator=configs.generator,458discriminator=configs.discriminator)

461withexperiment.start():

463configs.train()

467if\_\_name\_\_=='\_\_main\_\_':468main()