def __init__(self, config, dataloader):
self.config = config
# prepare input, label
self.dataloader = dataloader
with open('../../data/resized_celebA/gender_label.pkl', 'rb') as fp:
gender_label = pickle.load(fp)
self.gender_label = torch.LongTensor(gender_label).squeeze()
self.nz = config.nz
self.ngf = config.ngf
self.ndf = config.ndf
self.nch = config.nch
self.n_classes = config.n_classes
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_net()
self.setup_label()
self.setup_fix_vectors()
def __init__(self, n_epochs, lr):
self.n_epochs = n_epochs
self.lr = lr
self.datasets, self.dataloaders = get_ds_loaders()
self.model = CustomNet().to(device)
self.criterion = torch.nn.CrossEntropyLoss()
self.build_opt()
self.vis = Visualizer()
def __init__(self, lr, n_epochs):
self.criterion = nn.CrossEntropyLoss()
self.n_epochs = n_epochs
self.lr = lr
self.build_net()
self.build_opt()
self.datasets, self.dataloaders = get_ds_loaders()
self.vis = Visualizer()
def __init__(self, config, dataloader):
self.config = config
self.dataloader = dataloader
self.nz = config.nz
self.ngf = config.ngf
self.ndf = config.ndf
self.n_classes = config.n_classes
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_net()
def __init__(self):
# options
self.opts = opts.train_options().parse()
# training data
self.train_data = Train_Data(self.opts)
self.train_dataloader = self.train_data.train_loader
# device
self.device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
# model
self.model = model(self.opts, self.device)
# visualizer
self.train_vis = Visualizer(env='Training')
def __init__(self, config, dataloader):
self.config = config
self.dataloader = dataloader
self.nch = int(config.nch)
self.nz = int(config.nz)
self.ngf = int(config.ngf)
self.ndf = int(config.ndf)
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_model()
def __init__(self, train_loader, test_loader, config):
self.train_loader = train_loader
self.test_loader = test_loader
self.config = config
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.num_epochs = config.num_epochs
self.lr = config.lr
self.in_channel = config.in_channel
self.image_size = config.image_size
self.hidden_dim = config.hidden_dim
self.output_dim = config.output_dim
self.log_interval = config.log_interval
self.sample_interval = config.sample_interval
self.ckpt_interval = config.ckpt_interval
self.sample_folder = config.sample_folder
self.ckpt_folder = config.ckpt_folder
self.build_net()
self.vis = Visualizer()
def train(self):
opt = optim.Adam([self.output], lr=self.lr, betas=[0.5, 0.999])
vis = Visualizer()
print("Learning started!!!")
for epoch in range(self.nepochs):
output_features = self.vgg(self.output)
content_features = self.vgg(self.content)
style_features = self.vgg(self.style)
style_loss = 0
content_loss = 0
for f_out, f_content, f_style in zip(output_features,
content_features,
style_features):
content_loss += torch.mean((f_out - f_content)**2)
# reshape
_, c, h, w = f_out.size()
f_out = f_out.view(c, h * w)
f_style = f_style.view(c, h * w)
f_out = torch.mm(f_out, f_out.t())
f_style = torch.mm(f_style, f_style.t())
style_loss += torch.mean((f_out - f_style)**2) / (c * h * w)
loss = content_loss + self.config.style_weight * style_loss
opt.zero_grad()
loss.backward()
opt.step()
# do logging
if (epoch + 1) % self.log_interval == 0:
print(
"[{}/{}]: Content loss: {:.4f}, Style loss: {:.4f}".format(
epoch + 1, self.nepochs, content_loss.item(),
style_loss.item()))
vis.plot("Content loss per %d epochs" % self.log_interval,
content_loss.item())
vis.plot("Style loss per %d epochs" % self.log_interval,
style_loss.item())
if (epoch + 1) % self.sample_interval == 0:
denorm = transforms.Normalize((-2.12, -2.04, -1.80),
(4.37, 4.46, 4.44))
img = self.output.clone().squeeze()
img = denorm(img).clamp_(0, 1)
vutils.save_image(
img,
"%s\\output-%04d.png" % (self.config.out_folder, epoch))
print("Learning finished!!!")
class Trainer(object):
# initializer
def __init__(self, config, dataloader):
self.config = config
self.dataloader = dataloader
self.nz = config.nz
self.ngf = config.ngf
self.ndf = config.ndf
self.n_classes = config.n_classes
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_net()
# network init
def build_net(self):
self.g = Generator(self.nz, self.ngf, self.n_classes, self.image_size)
self.g.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.g != '':
self.g.load_state_dict(torch.load(self.config.g))
self.g = self.g.to(device)
self.d = Discriminator(self.ndf, self.n_classes, self.image_size)
self.d.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.d != '':
self.d.load_state_dict(torch.load(self.config.d))
self.d = self.d.to(device)
# transform label to onehot format
def get_onehot(self, label):
step_batch = label.size(0)
label = label.long().to(device)
oneHot = torch.zeros(step_batch, self.n_classes).to(device)
oneHot.scatter_(1, label.view(step_batch, 1), 1)
oneHot = oneHot.to(device)
return oneHot
def denorm(self, x):
out = (x + 1) / 2
return out.clamp(0, 1)
# training method
def train(self):
# setup loss function, optimizers
criterion = nn.BCELoss()
g_opt = optim.Adam(self.g.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
d_opt = optim.Adam(self.d.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
# setup fixed noise, label
fixed_noise = torch.FloatTensor(50, self.nz).normal_(0, 1).to(device)
fixed_label = (torch.rand(50, 1) * self.n_classes).long().to(device)
fixed_label = self.get_onehot(fixed_label)
print("Learning started!!")
for epoch in range(self.n_epochs):
# learning rate decay
if (epoch + 1) == 30:
g_opt.param_groups[0]['lr'] /= 10
d_opt.param_groups[0]['lr'] /= 10
print("Learning rate change!")
if (epoch + 1) == 40:
g_opt.param_groups[0]['lr'] /= 10
d_opt.param_groups[0]['lr'] /= 10
print("Learning rate change!")
for step, (x_real, y_) in enumerate(self.dataloader):
# ================================================================== #
# Train the discriminator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = True
step_batch = x_real.size(0)
x_real = x_real.view(-1, self.image_size**2).to(
device) # real data X: batch, 28*28
y_real = self.get_onehot(y_) # real data Y: batch, 10
target_real = torch.ones(step_batch).to(
device) # target for real: batch
target_fake = torch.zeros(step_batch).to(
device) # target tor fake: batch
# compute output and loss with real data
D_out_from_real = self.d(x_real, y_real).squeeze() # batch
D_x = D_out_from_real.data.mean()
D_loss_from_real = criterion(D_out_from_real, target_real)
# compute output and loss with fake data
random_z = torch.rand((step_batch, self.nz)).to(device)
random_y = (torch.rand(step_batch, 1) *
self.n_classes).long().to(device)
random_y = self.get_onehot(random_y)
x_fake = self.g(random_z, random_y).to(device) # batch, 28*28
D_out_from_fake = self.d(x_fake, random_y).squeeze()
D_G_z1 = D_out_from_fake.data.mean()
D_loss_from_fake = criterion(D_out_from_fake, target_fake)
D_loss = D_loss_from_real + D_loss_from_fake
# reset + forward + backward
self.d.zero_grad()
D_loss.backward()
d_opt.step()
# ================================================================== #
# Train the Generator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = False
random_z = torch.rand((step_batch, self.nz)).to(device)
random_y = (torch.rand(step_batch, 1) *
self.n_classes).long().to(device)
random_y = self.get_onehot(random_y)
x_fake = self.g(random_z, random_y).to(device) # batch, 28*28
D_out_from_fake = self.d(x_fake, random_y).squeeze()
D_G_z2 = D_out_from_fake.data.mean()
G_loss = criterion(D_out_from_fake, target_real)
# reset + forward + backward
self.g.zero_grad()
G_loss.backward()
g_opt.step()
if step % 100 == 0:
print(
"[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f"
% (epoch + 1, self.n_epochs, step + 1,
len(self.dataloader), D_loss.item(), G_loss.item(),
D_x, D_G_z1, D_G_z2))
# plot to visdom
self.vis.plot("D loss per 100 steps", D_loss.item())
self.vis.plot("G loss per 100 steps", G_loss.item())
# save results
x_real = x_real.view(-1, 1, self.image_size,
self.image_size)
x_real = self.denorm(x_real)
vutils.save_image(x_real,
'%s/real_samples.png' % self.out_folder)
fake = self.g(fixed_noise, fixed_label)
fake = fake.view(-1, 1, self.image_size, self.image_size)
fake = self.denorm(fake)
vutils.save_image(
fake, '%s/fake_samples_epoch_%03d.png' %
(self.out_folder, epoch))
if epoch % 10 == 0:
# save checkpoints
torch.save(self.g.state_dict(),
'%s/G_epoch_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(),
'%s/D_epoch_%03d.pth' % (self.out_folder, epoch))
print("learning finished!")
torch.save(self.g.state_dict(),
'%s/G_final_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(),
'%s/D_final_%03d.pth' % (self.out_folder, epoch))
print("save checkpoint finished!")
def train(self):
vis = Visualizer()
bce = nn.BCELoss()
cae = nn.L1Loss()
size_PatchGAN = 30
# setup optimizers
if self.config.rmsprop:
optG = optim.RMSprop(self.netG.parameters(), lr=self.lrG)
optD = optim.RMSprop(self.netD.parameters(), lr=self.lrD)
else:
optG = optim.Adam(self.netG.parameters(),
lr=self.lrG,
betas=(self.beta1, 0.999),
weight_decay=0.0)
optD = optim.Adam(self.netD.parameters(),
lr=self.lrD,
betas=(self.beta1, 0.999),
weight_decay=self.weight_decay)
# training loop
gan_iter = 0
start_time = time.time()
for epoch in range(self.nepochs):
for step, data in enumerate(self.dataloader, 0):
# set models train mode
self.netD.train()
self.netG.train()
# facades-> imageA: target, imageB: input (B2A)
# data = [A | B]
if self.mode == "B2A":
target, input = data
elif self.mode == "A2B":
input, target = data
step_batch = target.size(0)
input, target = input.to(device), target.to(device)
targetD_real = torch.ones(step_batch, 1, size_PatchGAN,
size_PatchGAN)
targetD_fake = torch.ones(step_batch, 1, size_PatchGAN,
size_PatchGAN)
targetD_real, targetD_fake = targetD_real.to(
device), targetD_fake.to(device)
#=============================================#
# Train discriminator #
#=============================================#
for param in self.netD.parameters():
param.requires_grad = True
self.netD.zero_grad()
outD_real = self.netD(torch.cat((target, input),
1)) # conditional GAN
errD_real = bce(outD_real, targetD_real)
D_x = outD_real.data.mean()
x_hat = self.netG(input)
fake = x_hat.detach()
outD_fake = self.netD(torch.cat([fake, input],
1)) # conditional GAN
errD_fake = bce(outD_fake, targetD_fake)
errD = (errD_real + errD_fake) * 0.5 # combined loss
errD.backward()
D_G_z1 = outD_fake.data.mean()
optD.step()
# =============================================#
# Train generator #
# =============================================#
for param in self.netD.parameters():
param.requires_grad = False
self.netG.zero_grad()
# compute L_L1
if self.lambdaIMG != 0:
errG_l1 = cae(x_hat, target) * self.lambdaIMG
# compute L_cGAN
outD_fake = self.netD(torch.cat((x_hat, input),
1)) # conditional
targetD_real = torch.ones(step_batch, 1, size_PatchGAN,
size_PatchGAN).to(device)
if self.lambdaGAN != 0:
errG_gan = bce(outD_fake, targetD_real)
D_G_z2 = outD_fake.data.mean()
# combined loss
errG = errG_l1 + errG_gan
errG.backward()
optG.step()
gan_iter += 1
if gan_iter % self.log_interval == 0:
end_time = time.time()
print(
"[%d/%d] [%d/%d] time:%f D loss:%.3f G_L1 loss:%.3f G_gan loss:%.3f D(x)=%.3f D(G(z))=%.3f/ %.3f"
% (epoch + 1, self.nepochs, step + 1,
len(self.dataloader), end_time - start_time,
errD.item(), errG_l1.item(), errG_gan.item(), D_x,
D_G_z1, D_G_z2))
sys.stdout.flush()
self.train_logger.write('%d\t%f\t%f\t%f\t%f\t%f\t%f\n' % \
(gan_iter, errD.item(), errG_l1.item(), errG_gan.item(), D_x, D_G_z1, D_G_z2))
self.train_logger.flush()
vis.plot("D loss per %d steps" % self.log_interval,
errD.item())
vis.plot("G_L1 loss per %d steps" % self.log_interval,
errG_l1.item())
vis.plot("G_gan loss per %d steps" % self.log_interval,
errG_gan.item())
# do checkpointing
# torch.save(self.netG.state_dict(), "%s/netG_epoch_%d.pth" % (self.out_folder, epoch+169))
# torch.save(self.netD.state_dict(), "%s/netD_epoch_%d.pth" % (self.out_folder, epoch+169))
# do validating
self.netD.eval()
self.netG.eval()
val_batch_output = torch.zeros(self.val_input.size())
for idx in range(self.val_input.size(0)):
img = self.val_input[idx, :, :, :].unsqueeze(0)
input_img = Variable(img, volatile=True).to(device)
x_hat_val = self.netG(input_img)
val_batch_output[idx, :, :, :].copy_(x_hat_val.data[0])
vutils.save_image(
val_batch_output,
"%s/generated_epoch%03d.png" % (self.out_folder, epoch + 169))
print("Learning finished!")
torch.save(self.netG.state_dict(),
"%s/netG_final.pth" % self.out_folder)
torch.save(self.netD.state_dict(),
"%s/netD_final.pth" % self.out_folder)
print("Saving PTH finished!")
class Trainer(object):
# initializer
def __init__(self, config, dataloader):
self.config = config
self.dataloader = dataloader
self.nch = int(config.nch)
self.nz = int(config.nz)
self.ngf = int(config.ngf)
self.ndf = int(config.ndf)
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_model()
# building network
def build_model(self):
self.g = dcgan.Generator(self.nz, self.ngf, self.nch)
self.g.apply(weights_init)
# if trained weights exists
if self.config.g != '':
self.g.load_state_dict(torch.load(self.config.g))
self.g.to(device)
self.d = dcgan.Discriminator(self.ndf, self.nch)
self.d.apply(weights_init)
# if trained weights exists
if self.config.d != '':
self.d.load_state_dict(torch.load(self.config.d))
self.d.to(device)
# trainer method
def train(self):
criterion = nn.BCELoss()
fixed_noise = torch.FloatTensor(self.batch_size, self.nz, 1,
1).normal_(0, 1).to(device)
# setup optimizers
g_opt = optim.Adam(self.g.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
d_opt = optim.Adam(self.d.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
print("Learning started!")
for epoch in range(self.n_epochs):
for step, (real_data, _) in enumerate(self.dataloader):
# ================================================================== #
# Train the discriminator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = True
real_data = real_data.to(device)
step_batch = real_data.size(0)
# create the labels
target_real = torch.ones(step_batch).to(device)
target_fake = torch.zeros(step_batch).to(device)
# train with real data
self.d.zero_grad()
D_out_from_real = self.d(real_data)
D_loss_from_real = criterion(D_out_from_real, target_real)
D_x = D_out_from_real.data.mean()
# train with fake data
z = torch.randn(step_batch, self.nz).to(device)
z = z.view(-1, self.nz, 1, 1)
fake_data = self.g(z)
D_out_from_fake = self.d(fake_data)
D_loss_from_fake = criterion(D_out_from_fake, target_fake)
D_G_z1 = D_out_from_fake.data.mean()
# loss + forward + backward
D_loss = D_loss_from_real + D_loss_from_fake
D_loss.backward()
d_opt.step()
# ================================================================== #
# Train the generator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = False
self.g.zero_grad()
z = torch.randn(step_batch, self.nz).to(device)
z = z.view(-1, self.nz, 1, 1)
fake_data = self.g(z)
D_out_from_fake = self.d(fake_data)
D_G_z2 = D_out_from_fake.data.mean()
# loss + forward + backward
G_loss = criterion(D_out_from_fake, target_real)
G_loss.backward()
g_opt.step()
if step % 100 == 0:
print(
'[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
% (epoch, self.n_epochs, step, len(self.dataloader),
D_loss.item(), G_loss.item(), D_x, D_G_z1, D_G_z2))
# plot to visdom
self.vis.plot("D loss per 100 step", D_loss.item())
self.vis.plot("G loss per 100 step", G_loss.item())
# save images
vutils.save_image(real_data,
'%s/real_samples.png' % self.out_folder,
normalize=True)
fake = self.g(fixed_noise)
vutils.save_image(fake.data,
'%s/fake_samples_epoch_%03d.png' %
(self.out_folder, epoch),
normalize=True)
if epoch % 10 == 0:
# save checkpoints
torch.save(self.g.state_dict(),
'%s/G_epoch_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(),
'%s/D_epoch_%03d.pth' % (self.out_folder, epoch))
print("learning finished!")
torch.save(self.g.state_dict(),
'%s/G_final_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(),
'%s/D_final_%03d.pth' % (self.out_folder, epoch))
print("save checkpoint finished!")
batch_size=batch_size, shuffle=True)
sample_dir = 'samples'
# Create a directory if not exists
if not os.path.exists(sample_dir):
os.makedirs(sample_dir)
def denorm(x):
out = (x + 1) / 2
return out.clamp(0, 1)
if __name__ == "__main__":
# set visdom tool
vis = Visualizer()
# build network
G = Generator(128).to(device)
D = Discriminator(128).to(device)
G.weight_init(mean=0.0, std=0.02)
D.weight_init(mean=0.0, std=0.02)
# define loss function and optimizers
criterion = nn.BCELoss()
G_opt = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
D_opt = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))
# reset gradient of optimizers
def reset_grad():
class Train():
def __init__(self, n_epochs, lr):
self.n_epochs = n_epochs
self.lr = lr
self.datasets, self.dataloaders = get_ds_loaders()
self.model = CustomNet().to(device)
self.criterion = torch.nn.CrossEntropyLoss()
self.build_opt()
self.vis = Visualizer()
# for name, param in self.model.named_parameters():
# print(name, param.requires_grad)
def build_opt(self):
self.optimizer = optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=self.lr,
momentum=0.9) # train only FC layer parameter
# print(self.optimizer.param_groups) # => only attached layers params
# Decay LR by a factor of 0.1 every 7 epochs
self.exp_lr_scheduler = lr_scheduler.StepLR(self.optimizer, step_size=7, gamma=0.1)
def train(self):
start = time.time()
best_model_wts = copy.deepcopy(self.model.state_dict())
best_acc = 0.0
for epoch in range(self.n_epochs):
print("Epoch [%d/%d]" % (epoch + 1, self.n_epochs))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
self.exp_lr_scheduler.step()
self.model.train() # Set model to training mode
else:
self.model.eval() # Set model to evaluate mode
epoch_loss = []
epoch_acc = []
# Iterate over data.
for step, (inputs, labels) in enumerate(self.dataloaders[phase]):
inputs = inputs.to(device)
labels = labels.to(device)
# zero the parameter gradients
self.optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
# for param in self.model.parameters():
# print(param[0], param.requires_grad)
# => if param.requires_grad == False, param doesn't change
outputs = self.model(inputs)
_, preds = torch.max(outputs, 1)
loss = self.criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
self.optimizer.step()
# statistics
step_loss = loss.item()
step_acc = float(torch.sum(preds == labels.data)) / len(preds)
if phase == 'train':
print("[%d/%d] [%d/%d] loss: %.3f acc: %.3f" % (
epoch + 1, self.n_epochs, step + 1, len(self.dataloaders[phase]), step_loss, step_acc))
self.vis.plot("Train loss plot per step", step_loss)
self.vis.plot("Train acc plot per step", step_acc)
epoch_loss.append(step_loss)
epoch_acc.append(step_acc)
epoch_loss = np.mean(epoch_loss)
epoch_acc = np.mean(epoch_acc)
print("[%d/%d] phase=%s: Avg loss: %.3f Avg acc: %.3f" % (
epoch + 1, self.n_epochs, phase, epoch_loss, epoch_acc))
self.vis.plot("%s avg loss plot per epoch" % phase, epoch_loss)
self.vis.plot("%s avg acc plot per epoch" % phase, epoch_acc)
# deep copy the model
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(self.model.state_dict())
print()
time_elapsed = time.time() - start
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
self.model.load_state_dict(best_model_wts)
return self.model
class Trainer(object):
def __init__(self, train_loader, test_loader, config):
self.train_loader = train_loader
self.test_loader = test_loader
self.config = config
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.num_epochs = config.num_epochs
self.lr = config.lr
self.in_channel = config.in_channel
self.image_size = config.image_size
self.hidden_dim = config.hidden_dim
self.output_dim = config.output_dim
self.log_interval = config.log_interval
self.sample_interval = config.sample_interval
self.ckpt_interval = config.ckpt_interval
self.sample_folder = config.sample_folder
self.ckpt_folder = config.ckpt_folder
self.build_net()
self.vis = Visualizer()
def build_net(self):
# define network
self.net = VAE(self.in_channel, self.image_size, self.hidden_dim,
self.output_dim)
if self.config.mode == 'test' and self.config.training_path == '':
print("[*] Enter model path!")
exit()
# if training model exists
if self.config.training_path != '':
self.net.load_state_dict(
torch.load(self.config.training_path,
map_location=lambda storage, loc: storage))
print("[*] Load weight from {}!".format(self.config.training_path))
self.net.to(self.device)
# define loss function
def loss_function(self, recon_x, x, mu, logvar):
criterion = nn.MSELoss(reduction='sum').to(self.device)
bce = criterion(recon_x,
x.view(-1, self.in_channel * (self.image_size**2)))
kld = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())
return bce + kld
def train(self):
# define optimizer
optimizer = Adam(self.net.parameters(), self.lr)
step = 0
print("[*] Learning started!")
# get fixed sample
temp_iter = iter(self.train_loader)
fixed_imgs, _ = next(temp_iter)
fixed_imgs = fixed_imgs.to(self.device)
# save fixed sample image
x_path = os.path.join(self.sample_folder, 'fixed_input.png')
save_image(fixed_imgs, x_path, normalize=True)
print("[*] Save fixed input image!")
# flatten
fixed_imgs = fixed_imgs.view(fixed_imgs.size(0), -1)
for epoch in range(self.num_epochs):
for i, (imgs, _) in enumerate(self.train_loader):
self.net.train()
imgs = imgs.view(imgs.size(0), -1)
imgs = imgs.to(self.device)
# forwarding and compute loss
recon, mu, logvar = self.net(imgs)
# testing code
# print("reconstructed:", recon.shape)
# print("original:", imgs.shape)
loss = self.loss_function(recon, imgs, mu, logvar)
# backwarding
optimizer.zero_grad()
loss.backward()
optimizer.step()
# do logging
if (step + 1) % self.log_interval == 0:
print("[{}/{}] [{}/{}] Loss:{:3f}".format(
epoch + 1, self.num_epochs, i + 1,
len(self.train_loader),
loss.item() / len(imgs)))
self.vis.plot("loss plot", loss.item() / len(imgs))
# do sampling
if (step + 1) % self.sample_interval == 0:
recon, mu, logvar = self.net(fixed_imgs)
recon = recon.view(-1, self.in_channel, self.image_size,
self.image_size)
x_hat_path = os.path.join(
self.sample_folder,
'output_epoch{}.png'.format(epoch + 1))
save_image(recon, x_hat_path, normalize=True)
print("[*] Save sample images!")
step += 1
if (epoch + 1) % self.ckpt_interval == 0:
ckpt_path = os.path.join(self.ckpt_folder,
'ckpt_epoch{}.pth'.format(epoch + 1))
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Checkpoint saved!")
print("[*] Learning finished!")
ckpt_path = os.path.join(self.ckpt_folder, 'final_model.pth')
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Final weight saved!")
class Trainer(object):
def __init__(self, train_loader, test_loader, config):
self.train_loader = train_loader
self.test_loader = test_loader
self.config = config
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.num_epochs = config.num_epochs
self.lr = config.lr
self.in_channel = config.in_channel
self.image_size = config.image_size
self.hidden_dim = config.hidden_dim
self.output_dim = config.output_dim
self.log_interval = config.log_interval
self.sample_interval = config.sample_interval
self.ckpt_interval = config.ckpt_interval
self.sample_folder = config.sample_folder
self.ckpt_folder = config.ckpt_folder
self.build_net()
self.vis = Visualizer()
def build_net(self):
# define network
self.net = AutoEncoder(self.in_channel, self.image_size,
self.hidden_dim, self.output_dim)
if self.config.mode == 'test' and self.config.training_path == '':
print("[*] Enter model path!")
exit()
# if training model exists
if self.config.training_path != '':
self.net.load_state_dict(
torch.load(self.config.training_path,
map_location=lambda storage, loc: storage))
print("[*] Load weight from {}!".format(self.config.training_path))
self.net.to(self.device)
# add noise to image
def add_noise(self, imgs):
noise = torch.randn(imgs.size()) * 0.4
noisy_imgs = noise + imgs
return noisy_imgs
def train(self):
# define loss function
bce_criterion = nn.BCELoss().to(self.device)
mse_criterion = nn.MSELoss().to(self.device)
# define optimizer
optimizer = Adam(self.net.parameters(), self.lr)
step = 0
print("[*] Learning started!")
# get fixed sample
temp_iter = iter(self.train_loader)
fixed_imgs, _ = next(temp_iter)
fixed_imgs = fixed_imgs.to(self.device)
# save fixed sample image
x_path = os.path.join(self.sample_folder, 'fixed_input.png')
save_image(fixed_imgs, x_path, normalize=True)
print("[*] Save fixed input image!")
# make fixed noisy sample and save
fixed_noisy_imgs = self.add_noise(fixed_imgs)
noisy_x_path = os.path.join(self.sample_folder,
'fixed_noisy_input.png')
save_image(fixed_noisy_imgs, noisy_x_path, normalize=True)
print("[*] Save fixed noisy input image!")
# flatten data tensors
fixed_imgs = fixed_imgs.view(fixed_imgs.size(0), -1)
fixed_noisy_imgs = fixed_noisy_imgs.view(fixed_imgs.size(0), -1)
for epoch in range(self.num_epochs):
for i, (imgs, _) in enumerate(self.train_loader):
self.net.train()
imgs = imgs.view(imgs.size(0), -1) # original images
noisy_imgs = self.add_noise(imgs) # add noise
noisy_imgs = noisy_imgs.to(self.device)
# forwarding
outputs = self.net(noisy_imgs) # use noisy image as input
bce_loss = bce_criterion(outputs, imgs)
mse_loss = mse_criterion(outputs, imgs)
# backwarding
optimizer.zero_grad()
bce_loss.backward() # backward BCE loss
optimizer.step()
# do logging
if (step + 1) % self.log_interval == 0:
print("[{}/{}] [{}/{}] BCE loss: {:3f}, MSE loss:{:3f}".
format(epoch + 1, self.num_epochs, i + 1,
len(self.train_loader),
bce_loss.item() / len(imgs),
mse_loss.item() / len(imgs)))
self.vis.plot("BCE Loss plot", bce_loss.item() / len(imgs))
self.vis.plot("MSE Loss plot", mse_loss.item() / len(imgs))
# do sampling
if (step + 1) % self.sample_interval == 0:
outputs = self.net(fixed_noisy_imgs)
x_hat = outputs.cpu().data.view(outputs.size(0), -1,
self.image_size,
self.image_size)
x_hat_path = os.path.join(
self.sample_folder,
'output_epoch{}.png'.format(epoch + 1))
save_image(x_hat, x_hat_path, normalize=True)
print("[*] Save sample images!")
step += 1
if (epoch + 1) % self.ckpt_interval == 0:
ckpt_path = os.path.join(self.ckpt_folder,
'ckpt_epoch{}.pth'.format(epoch + 1))
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Checkpoint saved!")
print("[*] Learning finished!")
ckpt_path = os.path.join(self.ckpt_folder, 'final_model.pth')
torch.save(self.net.state_dict(), ckpt_path)
print("[*] Final weight saved!")
from vis_tool import Visualizer
model = CNN().cuda()
# model.l
model.load_state_dict(torch.load('cnn.pkl'))
print(model)
_, _, test_loader = get_loader('../Dataset/HV', '../Dataset/RV',
'../Dataset/testRV')
test_avg_acc = 0.
test_cnt = 0
savelist = []
vis = Visualizer()
for idx, (video, label, filename) in enumerate(test_loader):
video = video[0]
label = label[0]
filename = filename[0]
start = 0
end = 20
out_acc = 0.
out_predicted = []
snp_cnt = 0
test_cnt += 1
class Trainer(object):
# initializer
def __init__(self, config, dataloader):
self.config = config
self.dataloader = dataloader
self.nz = config.nz
self.ngf = config.ngf
self.ndf = config.ndf
self.nch = config.nch
self.n_classes = config.n_classes
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_net()
# print(self.g)
# print(self.d)
# network init
def build_net(self):
self.g = Generator(self.nz, self.ngf, self.nch, self.n_classes)
self.g.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.g != '':
self.g.load_state_dict(torch.load(self.config.g))
self.g = self.g.to(device)
self.d = Discriminator(self.ndf, self.nch, self.n_classes)
self.d.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.d != '':
self.d.load_state_dict(torch.load(self.config.d))
self.d = self.d.to(device)
def denorm(self, x):
out = (x + 1) / 2
return out.clamp(0, 1)
# training method
def train(self):
# setup loss function, optimizers
criterion = nn.BCELoss()
g_opt = optim.Adam(self.g.parameters(), lr=self.lr, betas=(self.beta1, 0.999))
d_opt = optim.Adam(self.d.parameters(), lr=self.lr, betas=(self.beta1, 0.999))
# prepare label
onehot = torch.zeros(10, 10)
onehot = onehot.scatter_(1, torch.LongTensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]).view(10, 1), 1).view(10, 10, 1, 1)
fill = torch.zeros([10, 10, self.image_size, self.image_size])
for i in range(10):
fill[i, i, :, :] = 1 # shape: 10, 10, image_size, image_size
# setup fixed noise, label
fixed_noise = torch.FloatTensor(50, self.nz).normal_(0, 1).view(-1, self.nz, 1, 1).to(device)
fixed_label = (torch.rand(50, 1) * self.n_classes).long().to(device).squeeze().to(device)
fixed_onehot = onehot[fixed_label]
fixed_onehot = fixed_onehot.to(device)
print("Learning started!!")
for epoch in range(self.n_epochs):
# learning rate decay
if (epoch+1) == 11:
g_opt.param_groups[0]['lr'] /= 5
d_opt.param_groups[0]['lr'] /= 5
print("Learning rate change!")
if (epoch+1) == 16:
g_opt.param_groups[0]['lr'] /= 5
d_opt.param_groups[0]['lr'] /= 5
print("Learning rate change!")
for step, (x_real, y_) in enumerate(self.dataloader):
# ================================================================== #
# Train the discriminator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = True
step_batch = x_real.size(0)
target_real = torch.ones(step_batch).to(device)
target_fake = torch.zeros(step_batch).to(device)
x_real = x_real.to(device)
y_ = y_.to(device)
y_fill = fill[y_].to(device)
# compute output and loss with real data
D_out_from_real = self.d(x_real, y_fill)
D_x = D_out_from_real.data.mean()
D_loss_from_real = criterion(D_out_from_real, target_real)
z_ = torch.randn((step_batch, self.nz)).view(-1, self.nz, 1, 1).to(device)
y_ = (torch.rand(step_batch, 1) * self.n_classes).long().to(device).squeeze() # batch
y_label = onehot[y_].to(device) # batch, 10, 1, 1
y_fill = fill[y_].to(device) # batch, 10, 32, 32
x_fake = self.g(z_, y_label)
x_fake = x_fake.to(device)
# compute output and loss with fake data
D_out_from_fake = self.d(x_fake, y_fill)
D_G_z1 = D_out_from_fake.data.mean()
D_loss_from_fake = criterion(D_out_from_fake, target_fake)
D_loss = D_loss_from_real + D_loss_from_fake
# reset + forward + backward
self.d.zero_grad()
D_loss.backward()
d_opt.step()
# ================================================================== #
# Train the Generator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = False
z_ = torch.randn((step_batch, self.nz)).view(-1, self.nz, 1, 1).to(device)
y_ = (torch.rand(step_batch, 1) * self.n_classes).long().to(device).squeeze() # batch
y_label = onehot[y_].to(device) # batch, 10, 1, 1
y_fill = fill[y_].to(device) # batch, 10, 32, 32
x_fake = self.g(z_, y_label)
x_fake = x_fake.to(device)
D_out_from_fake = self.d(x_fake, y_fill)
D_G_z2 = D_out_from_fake.data.mean()
G_loss = criterion(D_out_from_fake, target_real)
# reset + forward + backward
self.g.zero_grad()
G_loss.backward()
g_opt.step()
if step % 100 == 0:
print("[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f"
% (epoch+1, self.n_epochs, step+1, len(self.dataloader),
D_loss.item(), G_loss.item(), D_x, D_G_z1, D_G_z2))
# plot to visdom
self.vis.plot("D loss per 100 steps", D_loss.item())
self.vis.plot("G loss per 100 steps", G_loss.item())
# save results
x_real = x_real.view(-1, 1, self.image_size, self.image_size)
x_real = self.denorm(x_real)
vutils.save_image(x_real,
'%s/real_samples.png' % self.out_folder)
fake = self.g(fixed_noise, fixed_onehot)
fake = fake.view(-1, 1, self.image_size, self.image_size)
fake = self.denorm(fake)
vutils.save_image(fake,
'%s/fake_samples_epoch_%03d.png' % (self.out_folder, epoch))
if epoch % 10 == 0:
# save checkpoints
torch.save(self.g.state_dict(), '%s/weights/G_epoch_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(), '%s/weights/D_epoch_%03d.pth' % (self.out_folder, epoch))
print("learning finished!")
torch.save(self.g.state_dict(), '%s/weights/G_final_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(), '%s/weights/D_final_%03d.pth' % (self.out_folder, epoch))
print("save checkpoint finished!")
class Trainer(object):
# initializer
def __init__(self, config, dataloader):
self.config = config
# prepare input, label
self.dataloader = dataloader
with open('../../data/resized_celebA/gender_label.pkl', 'rb') as fp:
gender_label = pickle.load(fp)
self.gender_label = torch.LongTensor(gender_label).squeeze()
self.nz = config.nz
self.ngf = config.ngf
self.ndf = config.ndf
self.nch = config.nch
self.n_classes = config.n_classes
self.batch_size = config.batch_size
self.image_size = config.image_size
self.lr = config.lr
self.beta1 = config.beta1
self.n_epochs = config.n_epochs
self.out_folder = config.out_folder
self.vis = Visualizer()
self.build_net()
self.setup_label()
self.setup_fix_vectors()
def denorm(self, x):
out = (x + 1) / 2
return out.clamp(0, 1)
# build network
def build_net(self):
self.g = Generator(self.nz, self.ngf, self.nch, self.n_classes)
self.g.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.g != '':
self.g.load_state_dict(torch.load(self.config.g))
self.g = self.g.to(device)
self.d = Discriminator(self.ndf, self.nch, self.n_classes)
self.d.weight_init(mean=0, std=0.02)
# if trained weight exists
if self.config.d != '':
self.d.load_state_dict(torch.load(self.config.d))
self.d = self.d.to(device)
# prepare label
def setup_label(self):
onehot = torch.zeros(self.n_classes, self.n_classes)
onehot = onehot.scatter_(1,
torch.LongTensor([0, 1]).view(2, 1),
1).view(2, 2, 1, 1)
fill = torch.zeros([2, 2, self.image_size, self.image_size])
for i in range(2):
fill[i, i, :, :] = 1
self.onehot = onehot
self.fill = fill
# setup fixed noise vectors
def setup_fix_vectors(self):
temp_z0_ = torch.randn(4, 100)
temp_z0_ = torch.cat([temp_z0_, temp_z0_], 0)
temp_z1_ = torch.randn(4, 100)
temp_z1_ = torch.cat([temp_z1_, temp_z1_], 0)
fixed_z_ = torch.cat([temp_z0_, temp_z1_], 0)
fixed_y_ = torch.cat(
[torch.zeros(4),
torch.ones(4),
torch.zeros(4),
torch.ones(4)], 0).type(torch.LongTensor).squeeze()
fixed_z_ = fixed_z_.view(-1, 100, 1, 1)
fixed_y_label_ = self.onehot[fixed_y_]
self.fixed_z = fixed_z_.to(device)
self.fixed_y_label = fixed_y_label_.to(device)
# save sample fake results
def save_sample_results(self, epoch, path, show=False, save=True):
self.g.eval()
fake_image = self.g(self.fixed_z, self.fixed_y_label)
self.g.train()
size_figure_grid = 4
fig, ax = plt.subplots(size_figure_grid,
size_figure_grid,
figsize=(5, 5))
for i, j in itertools.product(range(size_figure_grid),
range(size_figure_grid)):
ax[i, j].get_xaxis().set_visible(False)
ax[i, j].get_yaxis().set_visible(False)
for k in range(size_figure_grid * size_figure_grid):
i = k // size_figure_grid
j = k % size_figure_grid
ax[i, j].cla()
ax[i, j].imshow(
(fake_image[k].cpu().data.numpy().transpose(1, 2, 0) + 1) / 2)
label = 'Epoch {0}'.format(epoch)
fig.text(0.5, 0.04, label, ha='center')
if save:
plt.savefig(path)
if show:
plt.show()
else:
plt.close()
# training method
def train(self):
# Binary Cross Entropy loss
criterion = nn.BCELoss()
# Adam optimizer
g_opt = optim.Adam(self.g.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
d_opt = optim.Adam(self.d.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
print("Learning started!!")
for epoch in range(self.n_epochs):
# learning rate decay
if (epoch + 1) == 11:
g_opt.param_groups[0]['lr'] /= 5
d_opt.param_groups[0]['lr'] /= 5
print("Learning rate change!")
if (epoch + 1) == 16:
g_opt.param_groups[0]['lr'] /= 5
d_opt.param_groups[0]['lr'] /= 5
print("Learning rate change!")
for step, (x_real, _) in enumerate(self.dataloader):
# ================================================================== #
# Train the discriminator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = True
step_batch = x_real.size(0)
target_real = torch.ones(step_batch).to(device)
target_fake = torch.zeros(step_batch).to(device)
if step_batch != self.batch_size:
y_ = self.gender_label[self.batch_size * step:]
else:
y_ = self.gender_label[self.batch_size *
step:self.batch_size * (step + 1)]
x_real = x_real.to(device)
y_ = y_.to(device)
y_fill = self.fill[y_].to(device)
# compute output and loss with real data
D_out_from_real = self.d(x_real, y_fill)
D_x = D_out_from_real.data.mean()
D_loss_from_real = criterion(D_out_from_real, target_real)
z_ = torch.randn(
(step_batch, self.nz)).view(-1, self.nz, 1, 1).to(device)
y_ = (torch.rand(step_batch, 1) *
self.n_classes).long().to(device).squeeze() # batch
y_label = self.onehot[y_].to(device) # batch, 10, 1, 1
y_fill = self.fill[y_].to(device) # batch, 10, 32, 32
x_fake = self.g(z_, y_label)
x_fake = x_fake.to(device)
D_out_from_fake = self.d(x_fake, y_fill)
D_G_z1 = D_out_from_fake.data.mean()
D_loss_from_fake = criterion(D_out_from_fake, target_fake)
D_loss = D_loss_from_real + D_loss_from_fake
# reset + forward + backward
self.d.zero_grad()
D_loss.backward()
d_opt.step()
# ================================================================== #
# Train the Generator #
# ================================================================== #
for p in self.d.parameters():
p.requires_grad = False
z_ = torch.randn(
(step_batch, self.nz)).view(-1, self.nz, 1, 1).to(device)
y_ = (torch.rand(step_batch, 1) *
self.n_classes).long().to(device).squeeze() # batch
y_label = self.onehot[y_].to(device) # batch, 10, 1, 1
y_fill = self.fill[y_].to(device) # batch, 10, 32, 32
x_fake = self.g(z_, y_label)
x_fake = x_fake.to(device)
D_out_from_fake = self.d(x_fake, y_fill)
D_G_z2 = D_out_from_fake.data.mean()
G_loss = criterion(D_out_from_fake, target_real)
# reset + forward + backward
self.g.zero_grad()
G_loss.backward()
g_opt.step()
if step % 10 == 0:
# do logging
print(
"[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f"
% (epoch + 1, self.n_epochs, step + 1,
len(self.dataloader), D_loss.item(), G_loss.item(),
D_x, D_G_z1, D_G_z2))
# plot to visdom
self.vis.plot("D loss per 100 steps", D_loss.item())
self.vis.plot("G loss per 100 steps", G_loss.item())
# save results
x_real = x_real.view(-1, self.nch, self.image_size,
self.image_size)
vutils.save_image(self.denorm(x_real.data),
'%s/real_samples.png' % self.out_folder)
self.save_sample_results(
epoch, '%s/epoch%02d.png' % (self.out_folder, epoch))
print("Save result!!")
if epoch != 0 and epoch % 5 == 0:
os.makedirs('%s/weights' % self.out_folder, exist_ok=True)
# save checkpoints
torch.save(
self.g.state_dict(),
'%s/weights/G_epoch_%03d.pth' % (self.out_folder, epoch))
torch.save(
self.d.state_dict(),
'%s/weights/D_epoch_%03d.pth' % (self.out_folder, epoch))
print("learning finished!")
torch.save(self.g.state_dict(),
'%s/weights/G_final_%03d.pth' % (self.out_folder, epoch))
torch.save(self.d.state_dict(),
'%s/weights/D_final_%03d.pth' % (self.out_folder, epoch))
print("save checkpoint finished!")
def train(self):
vis = Visualizer()
optimizer = optim.Adam(self.tfNet.parameters(), self.lr, betas=[0.5, 0.999])
criterion = nn.MSELoss()
style = utils.load_image(self.style_path)
style = self.style_transform(style)
style = style.repeat(self.batch_size, 1, 1, 1).to(device)
style = utils.normalize_batch(style)
features_style = self.vggNet(style)
gram_style = [utils.gram_matrix(f) for f in features_style]
start_time = time.time()
print("Learning started!!!")
for epoch in range(self.nepochs):
for step, (content, _) in enumerate(self.dataloader):
self.tfNet.train()
step_batch = content.size(0)
optimizer.zero_grad()
content = content.to(device)
output = self.tfNet(content)
content = utils.normalize_batch(content)
output = utils.normalize_batch(output)
output_img = output
features_content = self.vggNet(content)
features_output = self.vggNet(output)
content_loss = self.content_weight * criterion(features_output.relu2_2,
features_content.relu2_2)
style_loss = 0.
for ft_output, gm_style in zip(features_output, gram_style):
gm_output = utils.gram_matrix(ft_output)
style_loss += criterion(gm_output,
gm_style[:step_batch, :, :])
style_loss *= self.style_weight
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
if (step+1) % self.log_interval == 0:
end_time = time.time()
print("[%d/%d] [%d/%d] time: %f content loss:%.4f style loss:%.4f total loss: %.4f"
% (epoch+1, self.nepochs, step+1, len(self.dataloader), end_time - start_time,
content_loss.item(), style_loss.item(), total_loss.item()))
vis.plot("Content loss per %d steps" % self.log_interval, content_loss.item())
vis.plot("Style loss per %d steps" % self.log_interval, style_loss.item())
vis.plot("Total loss per %d steps" % self.log_interval, total_loss.item())
# do save sample images
if (epoch+1) % self.sample_interval == 0:
img = output_img.cpu()
img = img[0]
utils.save_image("%s/output_epoch%d.png" % (self.sample_folder, epoch + 1), img)
# do checkpointing
if (epoch+1) % self.checkpoint_interval == 0:
self.tfNet.eval()
torch.save(self.tfNet.state_dict(), "%s/model_epoch%d.pth" % (self.sample_folder, epoch + 1))
print("Learning finished!!!")
self.tfNet.eval().cpu()
torch.save(self.tfNet.state_dict(), "%s/epoch%d_final.pth" % (self.sample_folder, self.nepochs))
print("Save model complete!!!")
def train(self):
vis = Visualizer()
# setup losses
criterion_GAN = nn.MSELoss()
criterion_cycle = nn.L1Loss()
criterion_identity = nn.L1Loss()
# setup optimizers
opt_G = optim.Adam(itertools.chain(self.netG_A2B.parameters(),
self.netG_B2A.parameters()),
lr=self.lr,
betas=(0.5, 0.999))
opt_D_A = optim.Adam(self.netD_A.parameters(),
lr=self.lr,
betas=(0.5, 0.999))
opt_D_B = optim.Adam(self.netD_B.parameters(),
lr=self.lr,
betas=(0.5, 0.999))
lambda_lr = LambdaLR(self.n_epochs, self.starting_epoch,
self.decay_epoch).step
# setup learning rate scheduler
lr_scheduler_G = optim.lr_scheduler.LambdaLR(opt_G,
lr_lambda=lambda_lr)
lr_scheduler_D_A = optim.lr_scheduler.LambdaLR(opt_D_A,
lr_lambda=lambda_lr)
lr_scheduler_D_B = optim.lr_scheduler.LambdaLR(opt_D_B,
lr_lambda=lambda_lr)
# setup buffers
fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()
print("Learning started!!!")
start_time = time.time()
for epoch in range(self.starting_epoch, self.n_epochs):
avg_loss_G = []
avg_loss_D = []
avg_loss_G_GAN = []
avg_loss_G_cycle = []
avg_loss_G_identity = []
for step, data in enumerate(self.dataloader):
self.netG_A2B.train()
self.netG_B2A.train()
real_A = data['A'].to(device)
real_B = data['B'].to(device)
# skip if image has 1 channel
if real_A.size(1) != 3 or real_B.size(1) != 3:
continue
step_batch = real_A.size(0)
target_real = torch.ones(step_batch,
requires_grad=False).to(device)
target_fake = torch.zeros(step_batch,
requires_grad=False).to(device)
# =============================================#
# Train Generator #
# =============================================#
for p in self.netD_A.parameters():
p.requires_grad = False
for p in self.netD_B.parameters():
p.requires_grad = False
opt_G.zero_grad()
# netG_A2B(B) should be equal to B if real B is fed
same_B = self.netG_A2B(real_B)
loss_identity_B = criterion_identity(same_B, real_B) * 5.0
# netG_B2A(A) should be equal to A if real A is fed
same_A = self.netG_B2A(real_A)
loss_identity_A = criterion_identity(same_A, real_A) * 5.0
# compute GAN loss
fake_B = self.netG_A2B(real_A)
outD_from_fake = self.netD_B(fake_B)
loss_GAN_A2B = criterion_GAN(outD_from_fake, target_real)
fake_A = self.netG_B2A(real_B)
outD_from_fake = self.netD_A(fake_A)
loss_GAN_B2A = criterion_GAN(outD_from_fake, target_real)
# compute Cycle loss
recovered_A = self.netG_B2A(fake_B)
loss_cycle_ABA = criterion_cycle(recovered_A, real_A) * 10.0
recovered_B = self.netG_A2B(fake_A)
loss_cycle_BAB = criterion_cycle(recovered_B, real_B) * 10.0
# compute Total loss
loss_G = loss_identity_A + loss_identity_B + loss_GAN_A2B + loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB
loss_G.backward()
opt_G.step()
# =============================================#
# Train Discriminator - A #
# =============================================#
for p in self.netD_A.parameters():
p.requires_grad = True
self.netD_A.train()
opt_D_A.zero_grad()
# compute real loss
outD_from_real = self.netD_A(real_A)
loss_D_from_real = criterion_GAN(outD_from_real, target_real)
# compute from fake
outD_from_fake = self.netD_A(fake_A.detach())
loss_D_from_fake = criterion_GAN(outD_from_fake, target_fake)
# compute total loss
loss_D_A = (loss_D_from_real + loss_D_from_fake) * 0.5
loss_D_A.backward()
opt_D_A.step()
# =============================================#
# Train Discriminator - B #
# =============================================#
for p in self.netD_B.parameters():
p.requires_grad = True
self.netD_B.train()
opt_D_B.zero_grad()
# compute real loss
outD_from_real = self.netD_B(real_B)
loss_D_from_real = criterion_GAN(outD_from_real, target_real)
# compute fake loss
outD_from_fake = self.netD_B(fake_B.detach())
loss_D_from_fake = criterion_GAN(outD_from_fake, target_fake)
# compute total loss
loss_D_B = (loss_D_from_real + loss_D_from_fake) * 0.5
loss_D_B.backward()
opt_D_B.step()
avg_loss_D.append((loss_D_A.item() + loss_D_B.item()) * 0.5)
avg_loss_G_identity.append(
(loss_identity_A.item() + loss_identity_B.item()) * 0.5)
avg_loss_G_cycle.append(
(loss_cycle_ABA.item() + loss_cycle_BAB.item()) * 0.5)
avg_loss_G_GAN.append(
(loss_GAN_A2B.item() + loss_GAN_B2A.item()) * 0.5)
avg_loss_G.append(loss_G.item())
if (step + 1) % self.log_interval == 0:
end_time = time.time()
print(
"[%d/%d] [%d/%d] time:%f loss_G:%.3f\tloss_G_identity:%.3f\tloss_G_GAN:%.3f\tloss_G_Cycle:%.3f\n"
"loss_D_A:%.3f\tloss_D_B:%.3f\t" %
(epoch + 1, self.n_epochs, step + 1,
len(self.dataloader), end_time - start_time, loss_G,
loss_identity_A + loss_identity_B,
loss_GAN_A2B + loss_GAN_B2A,
loss_cycle_ABA + loss_cycle_BAB, loss_D_A, loss_D_B))
print(
"========================================================================================="
)
lr_scheduler_G.step()
lr_scheduler_D_A.step()
lr_scheduler_D_B.step()
vis.plot("loss_G", np.mean(avg_loss_G))
vis.plot("loss_D", np.mean(avg_loss_D))
vis.plot("loss_G_GAN", np.mean(avg_loss_G_GAN))
vis.plot("loss_G_Cycle", np.mean(avg_loss_G_cycle))
vis.plot("loss_G_identity", np.mean(avg_loss_G_identity))
if (epoch + 1) % self.sample_interval == 0:
images = {
"real_A": real_A,
"real_B": real_B,
"fake_A": fake_A,
"fake_B": fake_B
}
self.sample_images(images)
print("Learning finished!!!")
torch.save(self.netG_A2B.state_dict(),
"%s\\netG_A2B.pth" % self.sample_folder)
torch.save(self.netG_B2A.state_dict(),
"%s\\netG_B2A.pth" % self.sample_folder)
torch.save(self.netD_A.state_dict(),
"%s\\netD_A.pth" % self.sample_folder)
torch.save(self.netD_B.state_dict(),
"%s\\netD_B.pth" % self.sample_folder)