class Model():
def __init__(self, args):
self.args = args
self.pretrained = False
self.epoch = 0
self.G = Generator()
self.D = Discriminator()
self.g_optimizer = optim.Adam(self.G.parameters(), lr=1E-4)
self.d_optimizer = optim.Adam(self.D.parameters(), lr=1E-4)
self.g_scheduler = optim.lr_scheduler.StepLR(self.g_optimizer,
step_size=40)
self.d_scheduler = optim.lr_scheduler.StepLR(self.d_optimizer,
step_size=40)
self.train_losses = []
self.val_losses = []
if args.load_model:
self._load_state(args.load_model)
# extract all layers prior to the last softmax of VGG-19
vgg19_layers = list(models.vgg19(pretrained=True).features)[:36]
self.vgg19 = nn.Sequential(*vgg19_layers).eval()
for param in self.vgg19.parameters():
param.requires_grad = False
self.mse_loss = torch.nn.MSELoss()
self.bce_loss = torch.nn.BCELoss()
def train(self, train_dataloader, val_dataloader=None):
self.D.to(device)
self.G.to(device)
self.vgg19.to(device)
""" Pretrain Generator """
if not self.pretrained:
log_message("Starting pretraining")
self._pretrain(train_dataloader)
self._save_state()
if val_dataloader:
val_g_loss, _ = self.evaluate(val_dataloader)
log_message("Pretrain G loss: {:.4f}".format(val_g_loss))
""" Real Training """
log_message("Starting training")
while self.epoch < self.args.epochs:
# Train one epoch
self.D.train()
self.G.train()
g_loss, d_loss = self._run_epoch(train_dataloader, train=True)
self.train_losses.append([g_loss, d_loss])
self.g_scheduler.step()
self.d_scheduler.step()
self.epoch += 1
log_message("Epoch: {}/{}".format(self.epoch, self.args.epochs))
# Print evaluation
train_string = "Train G loss: {:.4f} | Train D loss: {:.4f}".format(
g_loss, d_loss)
if self.epoch % self.args.eval_epochs == 0:
if val_dataloader:
val_g_loss, val_d_loss = self.evaluate(val_dataloader)
self.val_losses.append([val_g_loss, val_d_loss])
train_string += " | Val G loss: {:.4f} | Val D loss: {:.4f}".format(
val_g_loss, val_d_loss)
log_message(train_string)
# Save the model
if self.epoch % self.args.save_epochs == 0:
self._save_state()
log_message("Finished training")
self._save_state()
def evaluate(self, dataloader):
self.D.eval()
self.G.eval()
with torch.no_grad():
return self._run_epoch(dataloader, train=False)
def generate(self, dataloader):
def to_image(tensor):
array = tensor.data.cpu().numpy()
array = array.transpose((1, 2, 0))
array = np.clip(255.0 * (array + 1) / 2, 0, 255)
array = np.uint8(array)
return Image.fromarray(array)
self.D.eval()
self.G.eval()
if not os.path.exists(self.args.generate_dir):
os.mkdir(self.args.generate_dir)
with torch.no_grad():
for batch in dataloader:
low_res = batch['low_res'].to(device)
hi_res = batch['high_res']
generated = self.G(low_res)
for i in range(len(generated)):
naive = np.clip(
255.0 * low_res[i].data.cpu().numpy().transpose(
(1, 2, 0)), 0, 255)
naive = Image.fromarray(np.uint8(naive))
naive = naive.resize((96, 96), Image.BICUBIC)
fake_im = to_image(generated[i])
real_im = to_image(hi_res[i])
naive.save(
os.path.join(self.args.generate_dir,
"{}_naive.png".format(i)))
fake_im.save(
os.path.join(self.args.generate_dir,
"{}_fake.png".format(i)))
real_im.save(
os.path.join(self.args.generate_dir,
"{}_real.png".format(i)))
if i > 10:
return
def _load_state(self, fname):
if torch.cuda.is_available():
map_location = lambda storage, loc: storage.cuda()
else:
map_location = 'cpu'
state = torch.load(fname, map_location=map_location)
self.pretrained = state["pretrained"]
self.epoch = state["epoch"]
self.train_losses = state["train_losses"]
self.val_losses = state["val_losses"]
self.G.load_state_dict(state["G"])
self.D.load_state_dict(state["D"])
self.g_optimizer.load_state_dict(state["g_optimizer"])
self.d_optimizer.load_state_dict(state["d_optimizer"])
self.g_scheduler.load_state_dict(state["g_scheduler"])
self.d_scheduler.load_state_dict(state["d_scheduler"])
for state in self.d_optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.to(device)
for state in self.g_optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.to(device)
def _save_state(self):
if not os.path.exists(self.args.save_dir):
os.mkdir(self.args.save_dir)
fname = "%s/save_%d.pkl" % (self.args.save_dir, self.epoch)
state = {
"pretrained": self.pretrained,
"epoch": self.epoch,
"G": self.G.state_dict(),
"D": self.D.state_dict(),
"g_optimizer": self.g_optimizer.state_dict(),
"d_optimizer": self.d_optimizer.state_dict(),
"g_scheduler": self.g_scheduler.state_dict(),
"d_scheduler": self.d_scheduler.state_dict(),
"train_losses": self.train_losses,
"val_losses": self.val_losses
}
torch.save(state, fname)
def _pretrain(self, dataloader):
self.G.train()
for i in range(self.args.pretrain_epochs):
log_message("Pretrain Epoch: {}/{}".format(
i, self.args.pretrain_epochs))
for batch in dataloader:
low_res = batch['low_res'].to(device)
high_res = batch['high_res'].to(device)
self.g_optimizer.zero_grad()
generated = self.G(low_res)
# Optimize pixel loss
g_loss = self.mse_loss(generated, high_res)
g_loss.backward()
self.g_optimizer.step()
self.pretrained = True
def _run_epoch(self, dataloader, train):
g_losses, d_losses = [], []
for batch in dataloader:
low_res = batch['low_res'].to(device)
high_res = batch['high_res'].to(device)
batch_size = high_res.size(0)
real = torch.ones((batch_size, 1), requires_grad=False).to(device)
fake = torch.zeros((batch_size, 1), requires_grad=False).to(device)
""" Discriminator """
generated = self.G(low_res)
self.d_optimizer.zero_grad()
real_loss = self.bce_loss(self.D(high_res), real)
fake_loss = self.bce_loss(self.D(generated), fake)
d_loss = real_loss + fake_loss
d_losses.append(d_loss.item())
if train:
d_loss.backward()
self.d_optimizer.step()
""" Generator """
generated = self.G(low_res)
self.g_optimizer.zero_grad()
# take a [B, C, W, H] batch of [-1, 1] images, normalize, then run through vgg19
def vgg_features(image):
mean = torch.tensor(
[0.485, 0.456,
0.406]).unsqueeze(0).unsqueeze(2).unsqueeze(3).to(device)
std = torch.tensor(
[0.229, 0.224,
0.225]).unsqueeze(0).unsqueeze(2).unsqueeze(3).to(device)
image = (image + 1) / 2
image = (image - mean) / std
return self.vgg19(image)
pixel_loss = self.mse_loss(high_res, generated)
content_loss = self.mse_loss(vgg_features(high_res),
vgg_features(generated))
adversarial_loss = self.bce_loss(self.D(generated), real)
g_loss = pixel_loss + 0.006 * content_loss + 1E-3 * adversarial_loss
g_losses.append(g_loss.item())
if train:
g_loss.backward()
self.g_optimizer.step()
return np.mean(g_losses), np.mean(d_losses)
class Solver(object):
def __init__(self, config, data_loader):
self.generator = None
self.discriminator = None
self.g_optimizer = None
self.d_optimizer = None
self.g_conv_dim = config.g_conv_dim
self.d_conv_dim = config.d_conv_dim
self.z_dim = config.z_dim
self.beta1 = config.beta1
self.beta2 = config.beta2
self.image_size = config.image_size
self.data_loader = data_loader
self.num_epochs = config.num_epochs
self.batch_size = config.batch_size
self.sample_size = config.sample_size
self.lr = config.lr
self.log_step = config.log_step
self.sample_step = config.sample_step
self.sample_path = config.sample_path
self.model_path = config.model_path
self.epoch = config.epoch
self.build_model()
self.plotter = Plotter()
def build_model(self):
"""Build generator and discriminator."""
self.generator = Generator(z_dim=self.z_dim)
print(count_parameters(self.generator))
self.discriminator = Discriminator()
print(count_parameters(self.discriminator))
self.g_optimizer = optim.Adam(self.generator.parameters(),
self.lr, (self.beta1, self.beta2))
self.d_optimizer = optim.Adam(self.discriminator.parameters(),
self.lr*1, (self.beta1, self.beta2))
if self.epoch:
g_path = os.path.join(self.model_path, 'generator-%d.pkl' % self.epoch)
d_path = os.path.join(self.model_path, 'discriminator-%d.pkl' % self.epoch)
g_optim_path = os.path.join(self.model_path, 'gen-optim-%d.pkl' % self.epoch)
d_optim_path = os.path.join(self.model_path, 'dis-optim-%d.pkl' % self.epoch)
self.generator.load_state_dict(torch.load(g_path))
self.discriminator.load_state_dict(torch.load(d_path))
self.g_optimizer.load_state_dict(torch.load(g_optim_path))
self.d_optimizer.load_state_dict(torch.load(d_optim_path))
if torch.cuda.is_available():
self.generator.cuda()
self.discriminator.cuda()
def to_variable(self, x):
"""Convert tensor to variable."""
if torch.cuda.is_available():
x = x.cuda()
return Variable(x)
def to_data(self, x):
"""Convert variable to tensor."""
if torch.cuda.is_available():
x = x.cpu()
return x.data
def reset_grad(self):
"""Zero the gradient buffers."""
self.discriminator.zero_grad()
self.generator.zero_grad()
def denorm(self, x):
"""Convert range (-1, 1) to (0, 1)"""
out = (x + 1) / 2
return out.clamp(0, 1)
def train(self):
"""Train generator and discriminator."""
fixed_noise = self.to_variable(torch.randn(self.batch_size, self.z_dim))
total_step = len(self.data_loader)
for epoch in range(self.epoch, self.epoch + self.num_epochs) if self.epoch else range(self.num_epochs):
for i, images in enumerate(self.data_loader):
if len(images) != self.batch_size:
continue
# self.plotter.draw_kernels(self.discriminator)
for p in self.discriminator.parameters():
p.requires_grad = True
#===================== Train D =====================#
images = self.to_variable(images)
images.retain_grad()
batch_size = images.size(0)
noise = self.to_variable(torch.randn(batch_size, self.z_dim))
# Train D to recognize real images as real.
outputs = self.discriminator(images)
real_loss = torch.mean((outputs - 1) ** 2) # L2 loss instead of Binary cross entropy loss (this is optional for stable training)
# real_loss = torch.mean(outputs - 1)
# Train D to recognize fake images as fake.
fake_images = self.generator(noise)
fake_images.retain_grad()
outputs = self.discriminator(fake_images)
fake_loss = torch.mean(outputs ** 2)
# fake_loss = torch.mean(outputs)
# gradient penalty
gp_loss = calc_gradient_penalty(self.discriminator, images, fake_images)
# Backprop + optimize
d_loss = fake_loss + real_loss + gp_loss
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
if i % 10 == 0:
self.plotter.draw_activations(fake_images.grad[0], original=fake_images[0])
g_losses = []
for p in self.discriminator.parameters():
p.requires_grad = False
#===================== Train G =====================#
for g_batch in range(5):
noise = self.to_variable(torch.randn(batch_size, self.z_dim))
# Train G so that D recognizes G(z) as real.
fake_images = self.generator(noise)
outputs = self.discriminator(fake_images)
g_loss = torch.mean((outputs - 1) ** 2)
# g_loss = -torch.mean(outputs)
# Backprop + optimize
self.reset_grad()
g_loss.backward()
# if g_loss.item() < 0.5 * d_loss.item():
# break
self.g_optimizer.step()
g_losses.append("%.3f"%g_loss.clone().item())
# print the log info
if (i+1) % self.log_step == 0:
print('Epoch [%d/%d], Step[%d/%d], d_real_loss: %.4f, '
'd_fake_loss: %.4f, gp_loss: %s, g_loss: %s'
%(epoch+1, self.num_epochs, i+1, total_step,
real_loss.item(), fake_loss.item(), gp_loss.item(), ", ".join(g_losses)))
# save the sampled images
# print((i+1)%self.sample_step)
if (i) % self.sample_step == 0:
print("saving samples")
fake_images = self.generator(fixed_noise)
if not os.path.exists(self.sample_path):
os.makedirs(self.sample_path)
torchvision.utils.save_image(self.denorm(fake_images.data),
os.path.join(self.sample_path,
'fake_samples-%d-%d.png' %(epoch+1, i+1)))
# save the model parameters for each epoch
if epoch % 5 == 0:
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
g_path = os.path.join(self.model_path, 'generator-%d.pkl' %(epoch+1))
d_path = os.path.join(self.model_path, 'discriminator-%d.pkl' %(epoch+1))
g_optim_path = os.path.join(self.model_path, 'gen-optim-%d.pkl' % (epoch + 1))
d_optim_path = os.path.join(self.model_path, 'dis-optim-%d.pkl' % (epoch + 1))
torch.save(self.generator.state_dict(), g_path)
torch.save(self.discriminator.state_dict(), d_path)
torch.save(self.g_optimizer.state_dict(), g_optim_path)
torch.save(self.d_optimizer.state_dict(), d_optim_path)
def sample(self):
# Load trained parameters
g_path = os.path.join(self.model_path, 'generator-%d.pkl' % self.num_epochs)
d_path = os.path.join(self.model_path, 'discriminator-%d.pkl' % self.num_epochs)
self.generator.load_state_dict(torch.load(g_path))
self.discriminator.load_state_dict(torch.load(d_path))
self.generator.eval()
self.discriminator.eval()
# Sample the images
noise = self.to_variable(torch.randn(self.sample_size, self.z_dim))
fake_images = self.generator(noise)
sample_path = os.path.join(self.sample_path, 'fake_samples-final.png')
torchvision.utils.save_image(self.denorm(fake_images.data), sample_path, nrow=12)
print("Saved sampled images to '%s'" %sample_path)
class Solver(object):
def __init__(self, config, dataloader):
self.dataloader = dataloader
self.data_size = config.data_size
# self.iters = config.iters
self.loss_type = config.loss_type
self.G_lr = config.G_lr
self.D_lr = config.D_lr
self.beta1 = config.momentum
self.batch_size = config.batch_size
self.max_epoch = config.max_epoch
self.z_dim = config.z_dim
self.lr_update_step = config.lr_update_step
self.lr_decay_after = config.lr_decay_after
self.lr_decay = config.lr_decay
# path
self.sample_path = os.path.join(config.main_path, 'samples')
self.ckpt_path = os.path.join(config.main_path, 'checkpoints')
# misc
self.log_step = config.log_step
self.eval_step = config.eval_step
self.save_step = config.save_step
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.build_model()
def build_model(self):
self.G = Generator()
self.D = Discriminator()
self.G_optim = optim.Adam(self.G.parameters(), self.G_lr,
(self.beta1, 0.999))
self.D_optim = optim.Adam(self.D.parameters(), self.D_lr,
(self.beta1, 0.999))
if self.loss_type == 'BCEwL':
self.criterion = nn.BCEWithLogitsLoss()
elif self.loss_type == 'WGAN':
pass
elif self.loss_type == 'WGAN+':
pass
self.fixed_sample = None
self.fixed_noise = None
# self.true = torch.ones([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device)
# self.false = torch.zeros([self.batch_size, 1, 1, 1], requires_grad=False).to(self.device)
# Change to GPU mode
print('Change CPU mode to GPU mode...')
self.G.to(self.device)
self.D.to(self.device)
print('Creating models are success...')
def restore_model(self, resume_iters):
print('Load the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.ckpt_path, '{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.ckpt_path, '{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(torch.load(G_path))
self.D.load_state_dict(torch.load(D_path))
def train(self):
iters = self.max_epoch * len(self.dataloader)
data_iter = iter(self.dataloader)
self.fixed_sample = next(data_iter)
self.fixed_noise = torch.randn(self.batch_size,
self.z_dim).to(self.device)
num_data = 0
start_time = time.time()
print('Start training...')
for i in range(iters):
# try:
# sample = next(data_iter)
# except:
# print('error occur')
# data_iter = iter(self.dataloader)
# sample = next(data_iter)
sample = next(data_iter)
if i % len(self.dataloader) == 0:
data_iter = iter(self.dataloader)
# Load data.
right_embd = sample['right_embd'].to(self.device)
wrong_embd = sample['wrong_embd'].to(self.device)
z_noise = torch.randn(right_embd.size(0),
self.z_dim).to(self.device)
real_img = sample['real_img'].to(self.device)
fake_img = self.G(right_embd, z_noise)
# print('right_embd size: {}'.format(right_embd.size()))
# print('wrong_embd size: {}'.format(wrong_embd.size()))
# print('real_img size: {}'.format(real_img.size()))
num_data += right_embd.size(0)
T = torch.ones([right_embd.size(0), 1, 1, 1],
requires_grad=False).to(self.device)
F = torch.zeros([right_embd.size(0), 1, 1, 1],
requires_grad=False).to(self.device)
## Train Discriminator.
sr = self.D(real_img, right_embd) # {real image, right text}
rr_loss = self.criterion(sr, T)
sw = self.D(real_img, wrong_embd) # {real image, wrong text}
rw_loss = self.criterion(sw, F)
sf = self.D(fake_img.detach(),
right_embd) # {fake image, right text}
fr_loss = self.criterion(sf, F)
d_loss = rr_loss + rw_loss + fr_loss
## Backward and optimize for D.
self.D_optim.zero_grad()
d_loss.backward()
self.D_optim.step()
# For logs
loss = {}
loss['D/rr_loss'] = rr_loss.item()
loss['D/rw_loss'] = rw_loss.item()
loss['D/fr_loss'] = fr_loss.item()
loss['D/d_loss'] = d_loss.item()
## Train Generator.
sf = self.D(fake_img, right_embd)
g_loss = self.criterion(sf, T)
## Backward and optimize for G.
self.G_optim.zero_grad()
g_loss.backward()
self.G_optim.step()
loss['G/g_loss'] = g_loss.item()
## Print training information.
if (i + 1) % self.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
logs = "Elapsed [{}], Iter [{}/{}], Epoch [{}/{}]".format(
et, i + 1, iters, (i + 1) / len(self.dataloader),
self.max_epoch)
logs += ", Dataset [{}/{}]".format(num_data % self.data_size,
self.data_size)
for tag, value in loss.items():
logs += ', {} [{:.4f}]'.format(tag, value)
print(logs)
## Debug sample images.
if (i + 1) % self.eval_step == 0: #will be modified.
with torch.no_grad():
image_path = os.path.join(self.sample_path,
'{}.jpg'.format(i + 1))
fake_img = self.G(self.fixed_sample['right_embd'].to(
self.device), self.fixed_noise) #size: [B, 3, 64, 64]
real_img = self.fixed_sample['real_img']
img_list = []
for row in range(int(self.batch_size /
8)): #print multiple of 8 samples
img_list += [
real_img[row * 8 + col] for col in range(8)
]
img_list += [
fake_img[row * 8 + col].to('cpu')
for col in range(8)
]
sample_name = os.path.join(self.sample_path,
'{}iter.jpg'.format(i + 1))
save_image(make_grid(img_list), sample_name)
print('Save generated sample results {}iter.jpg into {}...'.
format(i + 1, self.sample_path))
## Save model checkpoints.
if (i + 1) % self.save_step == 0:
G_path = os.path.join(self.ckpt_path,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.ckpt_path,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Save model checkpoints into {}...'.format(
self.ckpt_path))
## Decay learning rates.
if (i + 1) % self.lr_update_step == 0:
if (i + 1) >= self.lr_decay_after:
self.G_lr = self.G_lr * self.lr_decay
self.D_lr = self.D_lr * self.lr_decay
for param_group in self.G_optim.param_groups:
param_group['lr'] = self.G_lr
for param_group in self.D_optim.param_groups:
param_group['lr'] = self.D_lr
print('Decay learning rates, g_lr: {}, d_lr: {}...'.format(
self.G_lr, self.D_lr))
def test(self):
pass
class Trainer(object):
def __init__(self, data_loader, config):
self.dataloader = data_loader
self.imsize = config.imsize
self.batch_size = config.batch_size
self.g_lr = config.g_lr
self.d_lr = config.d_lr
self.g_dim = config.g_dim
self.d_dim = config.d_dim
self.beta1 = config.beta1
self.beta2 = config.beta2
self.lambda_gp = config.lambda_gp
self.z_dim = config.z_dim
self.num_iters = config.total_step
self.num_iters_decay = config.iter_start_decay
self.log_step = config.log_step
self.sample_step = config.sample_step
self.lr_update_step = config.lr_iter_decay
self.lr_decay = config.lr_decay
self.model_save_step = config.model_save_step
self.resume_iters = config.resume_iter
self.version = config.version
self.device = torch.device('cuda:0')
self.sample_path = os.path.join(config.sample_path, self.version)
self.model_save_dir = os.path.join(config.model_save_path,
self.version)
self.build_model()
def build_model(self):
self.G = Generator(image_size=self.imsize,
z_dim=self.z_dim,
conv_dim=self.g_dim)
self.D = Discriminator(conv_dim=self.d_dim)
self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.print_network(self.G, 'G')
self.print_network(self.D, 'D')
self.G.to(self.device)
self.D.to(self.device)
def reset_grad(self):
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def update_lr(self, g_lr, d_lr):
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def print_network(self, model, name):
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(model)
print(name)
print("The number of parameters: {}".format(num_params))
def restore_model(self, resume_iters):
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def gradient_penalty(self, y, x):
weight = torch.ones(y.size()).to(self.device)
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm - 1)**2)
def train(self):
loss = {}
vis = visdom.Visdom()
data_iter = iter(self.dataloader)
g_lr = self.g_lr
d_lr = self.d_lr
fixed_z = torch.randn(self.batch_size, self.z_dim).cuda()
start_iters = 0
if self.resume_iters:
start_iters = self.resume_iters
self.restore_model(self.resume_iters)
print('start training')
start_time = time.time()
for i in range(start_iters, self.num_iters):
try:
x_mb, _ = next(data_iter)
except:
data_iter = iter(self.dataloader)
x_mb, _ = next(data_iter)
x_mb = x_mb.cuda()
z = torch.randn(x_mb.size(0), self.z_dim).cuda()
# train the discriminator
x_fake = self.G(z)
d_real = self.D(x_mb)
d_fake = self.D(x_fake)
d_loss_real = -torch.mean(d_real)
d_loss_fake = torch.mean(d_fake)
alpha = torch.rand(x_mb.size(0), 1, 1, 1).to(self.device)
# interpolate between real data and fake data
x_hat = (alpha * x_mb.data +
(1 - alpha) * x_fake.data).requires_grad_(True)
out_src = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
d_loss = d_loss_real + d_loss_fake + self.lambda_gp * d_loss_gp
d_loss.backward()
self.d_optimizer.step()
self.reset_grad()
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_gp'] = d_loss_gp.item()
# train generator
d_fake = self.D(self.G(z))
g_loss = -torch.mean(d_fake)
g_loss.backward()
self.g_optimizer.step()
self.reset_grad()
loss['G/loss'] = g_loss.item()
if (i + 1) % self.log_step == 0:
# visualize real and fake imgs
vis.images((x_fake + 1) / 2, win='fake_imgs')
vis.images((x_mb + 1) / 2, win='real_imgs')
# print and visualize losses
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Iteration [{}/{}]".format(
et, i + 1, self.num_iters)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
opts = dict(title='Losses',
width=13,
height=10,
legend=list(loss.keys()))
vis.line(Y=[list(loss.values())], X=[np.ones(len(loss))*(i+1)], win='Losses', \
update='append', opts=opts)
print(log)
if (i + 1) % self.lr_update_step == 0 and (
i + 1) > self.num_iters_decay:
g_lr = self.g_lr * self.lr_decay
d_lr = self.d_lr * self.lr_decay
self.update_lr(g_lr, d_lr)
print('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(
g_lr, d_lr))
# Sample images
if (i + 1) % self.sample_step == 0:
fake_images = self.G(fixed_z)
save_image(
denorm(fake_images.data),
os.path.join(self.sample_path,
'{}_fake.png'.format(i + 1)))
if (i + 1) % self.model_save_step == 0:
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(i + 1))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(i + 1))
torch.save(self.G.state_dict(), G_path)
torch.save(self.D.state_dict(), D_path)
print('Saved model checkpoints into {}...'.format(
self.model_save_dir))
def train(rank: int, cfg: DictConfig):
print(OmegaConf.to_yaml(cfg))
if cfg.train.n_gpu > 1:
init_process_group(backend=cfg.train.dist_config['dist_backend'],
init_method=cfg.train.dist_config['dist_url'],
world_size=cfg.train.dist_config['world_size'] *
cfg.train.n_gpu,
rank=rank)
device = torch.device(
'cuda:{:d}'.format(rank) if torch.cuda.is_available() else 'cpu')
generator = Generator(sum(cfg.model.feature_dims), *cfg.model.cond_dims,
**cfg.model.generator).to(device)
discriminator = Discriminator(**cfg.model.discriminator).to(device)
if rank == 0:
print(generator)
os.makedirs(cfg.train.ckpt_dir, exist_ok=True)
print("checkpoints directory : ", cfg.train.ckpt_dir)
if os.path.isdir(cfg.train.ckpt_dir):
cp_g = scan_checkpoint(cfg.train.ckpt_dir, 'g_')
cp_do = scan_checkpoint(cfg.train.ckpt_dir, 'd_')
steps = 1
if cp_g is None or cp_do is None:
state_dict_do = None
last_epoch = -1
else:
state_dict_g = load_checkpoint(cp_g, device)
state_dict_do = load_checkpoint(cp_do, device)
generator.load_state_dict(state_dict_g['generator'])
discriminator.load_state_dict(state_dict_do['discriminator'])
steps = state_dict_do['steps'] + 1
last_epoch = state_dict_do['epoch']
if cfg.train.n_gpu > 1:
generator = DistributedDataParallel(generator,
device_ids=[rank]).to(device)
discriminator = DistributedDataParallel(discriminator,
device_ids=[rank]).to(device)
optim_g = RAdam(generator.parameters(), cfg.opt.lr, betas=cfg.opt.betas)
optim_d = RAdam(discriminator.parameters(),
cfg.opt.lr,
betas=cfg.opt.betas)
if state_dict_do is not None:
optim_g.load_state_dict(state_dict_do['optim_g'])
optim_d.load_state_dict(state_dict_do['optim_d'])
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
optim_g, gamma=cfg.opt.lr_decay, last_epoch=last_epoch)
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(
optim_d, gamma=cfg.opt.lr_decay, last_epoch=last_epoch)
train_filelist = load_dataset_filelist(cfg.dataset.train_list)
trainset = FeatureDataset(cfg.dataset, train_filelist, cfg.data)
train_sampler = DistributedSampler(
trainset) if cfg.train.n_gpu > 1 else None
train_loader = DataLoader(trainset,
batch_size=cfg.train.batch_size,
num_workers=cfg.train.num_workers,
shuffle=True,
sampler=train_sampler,
pin_memory=True,
drop_last=True)
if rank == 0:
val_filelist = load_dataset_filelist(cfg.dataset.test_list)
valset = FeatureDataset(cfg.dataset,
val_filelist,
cfg.data,
segmented=False)
val_loader = DataLoader(valset,
batch_size=1,
num_workers=cfg.train.num_workers,
shuffle=False,
sampler=train_sampler,
pin_memory=True)
sw = SummaryWriter(os.path.join(cfg.train.ckpt_dir, 'logs'))
generator.train()
discriminator.train()
for epoch in range(max(0, last_epoch), cfg.train.epochs):
if rank == 0:
start = time.time()
print("Epoch: {}".format(epoch + 1))
if cfg.train.n_gpu > 1:
train_sampler.set_epoch(epoch)
for y, x_noised_features, x_noised_cond in train_loader:
if rank == 0:
start_b = time.time()
y = y.to(device, non_blocking=True)
x_noised_features = x_noised_features.transpose(1, 2).to(
device, non_blocking=True)
x_noised_cond = x_noised_cond.to(device, non_blocking=True)
z1 = torch.randn(cfg.train.batch_size,
cfg.model.cond_dims[1],
device=device)
z2 = torch.randn(cfg.train.batch_size,
cfg.model.cond_dims[1],
device=device)
y_hat1 = generator(x_noised_features, x_noised_cond, z=z1)
y_hat2 = generator(x_noised_features, x_noised_cond, z=z2)
# Discriminator
real_scores, fake_scores = discriminator(y), discriminator(
y_hat1.detach())
d_loss = discriminator_loss(real_scores, fake_scores)
optim_d.zero_grad()
d_loss.backward(retain_graph=True)
optim_d.step()
# Generator
g_stft_loss = criterion(y, y_hat1) + criterion(
y, y_hat2) - criterion(y_hat1, y_hat2)
g_adv_loss = adversarial_loss(fake_scores)
g_loss = g_adv_loss + g_stft_loss
optim_g.zero_grad()
g_loss.backward()
optim_g.step()
if rank == 0:
# STDOUT logging
if steps % cfg.train.stdout_interval == 0:
with torch.no_grad():
print(
'Steps : {:d}, Gen Loss Total : {:4.3f}, STFT Error : {:4.3f}, s/b : {:4.3f}'
.format(steps, g_loss, g_stft_loss,
time.time() - start_b))
# checkpointing
if steps % cfg.train.checkpoint_interval == 0:
ckpt_dir = "{}/g_{:08d}".format(cfg.train.ckpt_dir, steps)
save_checkpoint(
ckpt_dir, {
'generator':
(generator.module if cfg.train.n_gpu > 1 else
generator).state_dict()
})
ckpt_dir = "{}/do_{:08d}".format(cfg.train.ckpt_dir, steps)
save_checkpoint(
ckpt_dir, {
'discriminator':
(discriminator.module if cfg.train.n_gpu > 1 else
discriminator).state_dict(),
'optim_g':
optim_g.state_dict(),
'optim_d':
optim_d.state_dict(),
'steps':
steps,
'epoch':
epoch
})
# Tensorboard summary logging
if steps % cfg.train.summary_interval == 0:
sw.add_scalar("training/gen_loss_total", g_loss, steps)
sw.add_scalar("training/gen_stft_error", g_stft_loss,
steps)
# Validation
if steps % cfg.train.validation_interval == 0:
generator.eval()
torch.cuda.empty_cache()
val_err_tot = 0
with torch.no_grad():
for j, (y, x_noised_features,
x_noised_cond) in enumerate(val_loader):
y_hat = generator(
x_noised_features.transpose(1, 2).to(device),
x_noised_cond.to(device))
val_err_tot += criterion(y, y_hat).item()
if j <= 4:
# sw.add_audio('noised/y_noised_{}'.format(j), y_noised[0], steps, cfg.data.target_sample_rate)
sw.add_audio('generated/y_hat_{}'.format(j),
y_hat[0], steps,
cfg.data.sample_rate)
sw.add_audio('gt/y_{}'.format(j), y[0], steps,
cfg.data.sample_rate)
val_err = val_err_tot / (j + 1)
sw.add_scalar("validation/stft_error", val_err, steps)
generator.train()
steps += 1
scheduler_g.step()
scheduler_d.step()
if rank == 0:
print('Time taken for epoch {} is {} sec\n'.format(
epoch + 1, int(time.time() - start)))