| import torch
|
| import torch.nn.functional as F
|
| from IPython.display import clear_output
|
|
|
|
|
| class Trainer():
|
| def __init__(self, encoder, decoder, D, vgg, losses, data_len, ema=3, a_disc=1, a_vae=1, a_KL=0.1, isViT=True):
|
| self.vgg_schedule = None
|
| self.ema = 2/(ema+1)
|
| self.a_disc = a_disc
|
| self.a_vae = a_vae
|
| self.a_KL = a_KL
|
|
|
| self.isViT = isViT
|
| self.encoder = encoder
|
| self.decoder = decoder
|
| self.D = D
|
| self.vgg = vgg
|
| self.encoder_optimizer = torch.optim.Adam(self.encoder.parameters(), lr=1e-5)
|
| self.encoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.encoder_optimizer, T_max=50)
|
| self.decoder_optimizer = torch.optim.Adam(self.decoder.parameters(), lr=1e-5)
|
| self.decoder_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.decoder_optimizer, T_max=50)
|
| self.D_optimizer = torch.optim.Adam(self.D.parameters(), lr=4e-5)
|
| self.D_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.D_optimizer, T_max=50)
|
| self.losses = losses
|
| self.loss_vals = {loss:0 for loss in losses}
|
| self.data_len = data_len
|
| self.loss_record = []
|
| self.epoch = 1
|
| self.index = 1
|
| self.device = torch.device("cuda")
|
|
|
| self.encoder.to(self.device)
|
| self.decoder.to(self.device)
|
| self.D.to(self.device)
|
| self.vgg.to(self.device)
|
|
|
| def train_step(self, x, with_mse=False, freeze_ae=False, freeze_disc=False):
|
| self.index += 1
|
| x = x.to(self.device)
|
| with torch.no_grad():
|
| x_hat = self.decoder(self.encoder(x.permute(0,2,3,1))).permute(0,3,1,2) if not self.isViT else self.decoder(self.encoder(x))
|
| if not freeze_disc:
|
| disc_loss = F.relu(1. - self.D(x)).mean() + F.relu(1. + self.D(x_hat)).mean()
|
| self.D_optimizer.zero_grad()
|
| disc_loss.backward()
|
| self.D_optimizer.step()
|
| self.D_scheduler.step()
|
|
|
| if not freeze_ae:
|
| z = self.encoder(x.permute(0,2,3,1)) if not self.isViT else self.encoder(x)
|
| x_hat = self.decoder(z).permute(0,3,1,2) if not self.isViT else self.decoder(z)
|
| mse = F.mse_loss(x_hat, x)
|
| KL = 0.5 * (z.mean() ** 2)
|
| vgg_real = self.vgg(x)
|
| vgg_fake = self.vgg(x_hat)
|
| vgg_loss = 0
|
| for i in range(len(vgg_real)):
|
| vgg_loss += F.mse_loss(vgg_real[i], vgg_fake[i])
|
|
|
| adv_loss = 0
|
| if not freeze_disc:
|
| adv_loss = -(self.D(self.decoder(self.encoder(x))).mean())
|
|
|
| loss = mse * with_mse + self.a_KL* KL + vgg_loss + self.a_vae * adv_loss
|
| self.encoder_optimizer.zero_grad()
|
| self.decoder_optimizer.zero_grad()
|
| loss.backward()
|
| self.encoder_optimizer.step()
|
| self.decoder_optimizer.step()
|
| self.encoder_scheduler.step()
|
| self.decoder_scheduler.step()
|
|
|
| self.update_batch({"mse":mse.item() if not freeze_ae else 0,
|
| "gan":disc_loss.item() if not freeze_disc else 0,
|
| "vgg":vgg_loss.item() if not freeze_ae else 0,
|
| "KL":z.mean() if not freeze_ae else 0})
|
|
|
| def update_batch(self, loss_vals):
|
| clear_output(wait=True)
|
| for record in self.loss_record:
|
| print(record)
|
| self.loss_vals = {loss:(1-self.ema)*self.loss_vals[loss] + self.ema*loss_vals[loss] for loss in self.losses}
|
| print(f"epoch:{self.epoch} ", end="")
|
| for loss in self.losses:
|
| print(f"{loss}: {self.loss_vals[loss]:.3f} ", end="")
|
| for _ in range(int(self.index * 20 / self.data_len)):
|
| print("=", end="")
|
| for _ in range(int(self.index * 20 / self.data_len),20):
|
| print("-", end="")
|
|
|
| def update_epoch(self):
|
| self.index = 0
|
| record = f"epoch:{self.epoch} "
|
| for loss in self.losses:
|
| record += f"{loss}: {self.loss_vals[loss]:.3f} "
|
| self.loss_record.append(record)
|
| self.epoch += 1 |
