def main(opts):
# Create the data loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((opts.resolution, opts.resolution)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
drop_last=True
)
# Create the model
start_epoch = 0
G = G_stylegan2(fmap_base=opts.fmap_base,
resolution=opts.resolution,
mapping_layers=opts.mapping_layers,
opts=opts,
return_dlatents=True)
D = D_stylegan2(fmap_base=opts.fmap_base,
resolution=opts.resolution,
structure='resnet')
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO("Pre-trained weight cannot load successfully, train from scratch!")
# Multi-GPU support
if torch.cuda.device_count() > 1:
INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
G = torch.nn.DataParallel(G)
D = torch.nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# Create the criterion, optimizer and scheduler
lr_D = 0.0015
lr_G = 0.0015
optim_D = torch.optim.Adam(D.parameters(), lr=lr_D, betas=(0.9, 0.999))
# g_mapping has 100x lower learning rate
params_G = [{"params": G.g_synthesis.parameters()},
{"params": G.g_mapping.parameters(), "lr": lr_G * 0.01}]
optim_G = torch.optim.Adam(params_G, lr=lr_G, betas=(0.9, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = torch.nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img,) in enumerate(bar):
real_img = real_img.to(opts.device)
latents = torch.randn([real_img.size(0), 512]).to(opts.device)
# =======================================================================================================
# (1) Update D network: D_logistic_r1(default)
# =======================================================================================================
# Compute adversarial loss toward discriminator
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img, fake_dlatent = G(latents)
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit)
d_loss = d_loss + softplus(-real_logit)
# original
r1_penalty = D_logistic_r1(real_img.detach(), D)
d_loss = (d_loss + r1_penalty).mean()
# lite
# d_loss = d_loss.mean()
loss_D_list.append(d_loss.mean().item())
# Update discriminator
optim_D.zero_grad()
d_loss.backward()
optim_D.step()
# =======================================================================================================
# (2) Update G network: G_logistic_ns_pathreg(default)
# =======================================================================================================
# if i % CRITIC_ITER == 0:
G.zero_grad()
fake_scores_out = D(fake_img)
_g_loss = softplus(-fake_scores_out)
# Compute |J*y|.
# pl_noise = (torch.randn(fake_img.shape) / np.sqrt(fake_img.shape[2] * fake_img.shape[3])).to(fake_img.device)
# pl_grads = grad(torch.sum(fake_img * pl_noise), fake_dlatent, retain_graph=True)[0]
# pl_lengths = torch.sqrt(torch.sum(torch.sum(torch.mul(pl_grads, pl_grads), dim=2), dim=1))
# pl_mean = PL_DECAY * torch.sum(pl_lengths)
#
# pl_penalty = torch.mul(pl_lengths - pl_mean, pl_lengths - pl_mean)
# reg = pl_penalty * PL_WEIGHT
#
# # original
# g_loss = (_g_loss + reg).mean()
# lite
g_loss = _g_loss.mean()
loss_G_list.append(g_loss.mean().item())
# Update generator
g_loss.backward(retain_graph=True)
optim_G.step()
# Output training stats
bar.set_description(
"Epoch {} [{}, {}] [G]: {} [D]: {}".format(ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
with torch.no_grad():
fake_img = G(fix_z)[0].detach().cpu()
save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True)
# Save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def main(opts):
# Create the data loader
loader = sunnerData.DataLoader(
sunnerData.ImageDataset(root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((1024, 1024)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(
sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
)
# Create the model
start_epoch = 0
G = StyleGenerator()
D = StyleDiscriminator()
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO(
"Pre-trained weight cannot load successfully, train from scratch!")
# Multi-GPU support
if torch.cuda.device_count() > 1:
INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
G = nn.DataParallel(G)
D = nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# Create the criterion, optimizer and scheduler
optim_D = optim.Adam(D.parameters(), lr=0.00001, betas=(0.5, 0.999))
optim_G = optim.Adam(G.parameters(), lr=0.00001, betas=(0.5, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img, ) in enumerate(bar):
# =======================================================================================================
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
# =======================================================================================================
# Compute adversarial loss toward discriminator
D.zero_grad()
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device))
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit).mean()
d_loss = d_loss + softplus(-real_logit).mean()
if opts.r1_gamma != 0.0:
r1_penalty = R1Penalty(real_img.detach(), D)
d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5)
if opts.r2_gamma != 0.0:
r2_penalty = R2Penalty(fake_img.detach(), D)
d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5)
loss_D_list.append(d_loss.item())
# Update discriminator
d_loss.backward()
optim_D.step()
# =======================================================================================================
# (2) Update G network: maximize log(D(G(z)))
# =======================================================================================================
if i % CRITIC_ITER == 0:
G.zero_grad()
fake_logit = D(fake_img)
g_loss = softplus(-fake_logit).mean()
loss_G_list.append(g_loss.item())
# Update generator
g_loss.backward()
optim_G.step()
# Output training stats
bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format(
ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
with torch.no_grad():
fake_img = G(fix_z).detach().cpu()
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'),
nrow=4,
normalize=True)
# Save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def main(opts):
# Create the data loader
# loader = sunnerData.DataLoader(sunnerData.ImageDataset(
# root=[[opts.path]],
# transforms=transforms.Compose([
# sunnertransforms.Resize((1024, 1024)),
# sunnertransforms.ToTensor(),
# sunnertransforms.ToFloat(),
# #sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
# sunnertransforms.Normalize(),
# ])),
# batch_size=opts.batch_size,
# shuffle=True,
# )
loader = data_loader(opts.path)
device = fluid.CUDAPlace(0) if opts.device == 'GPU' else fluid.CPUPlace(0)
with fluid.dygraph.guard(device):
# Create the model
start_epoch = 0
G = StyleGenerator()
D = StyleDiscriminator()
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
#state = fluid.dygraph.load_dygraph(opts.resume)
state = load_checkpoint(opts.resume)
G.load_dict(state['G'])
D.load_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO(
"Pre-trained weight cannot load successfully, train from scratch!"
)
# # Multi-GPU support
# if torch.cuda.device_count() > 1:
# INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
# G = nn.DataParallel(G)
# D = nn.DataParallel(D)
scheduler_D = exponential_decay(learning_rate=0.00001,
decay_steps=1000,
decay_rate=0.99)
scheduler_G = exponential_decay(learning_rate=0.00001,
decay_steps=1000,
decay_rate=0.99)
optim_D = optim.Adam(parameter_list=D.parameters(),
learning_rate=scheduler_D)
optim_G = optim.Adam(parameter_list=G.parameters(),
learning_rate=scheduler_G)
# Train
fix_z = np.random.randn(opts.batch_size, 512)
fix_z = dygraph.to_variable(fix_z)
softplus = SoftPlus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
D.train()
G.train()
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader())
loss_D_list = []
loss_G_list = []
for i, data in enumerate(bar):
# =======================================================================================================
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
# =======================================================================================================
# Compute adversarial loss toward discriminator
real_img = np.array([item for item in data],
dtype='float32').reshape(
(-1, 3, 1024, 1024))
D.clear_gradients()
real_img = dygraph.to_variable(real_img)
real_logit = D(real_img)
z = np.float32(np.random.randn(real_img.shape[0], 512))
fake_img = G(dygraph.to_variable(z))
fake_logit = D(fake_img)
d_loss = layers.mean(softplus(fake_logit))
d_loss = d_loss + layers.mean(softplus(-real_logit))
if opts.r1_gamma != 0.0:
r1_penalty = R1Penalty(real_img, D)
d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5)
if opts.r2_gamma != 0.0:
r2_penalty = R2Penalty(fake_img, D)
d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5)
loss_D_list.append(d_loss.numpy())
# Update discriminator
d_loss.backward()
optim_D.minimize(d_loss)
# =======================================================================================================
# (2) Update G network: maximize log(D(G(z)))
# =======================================================================================================
if i % CRITIC_ITER == 0:
G.clear_gradients()
fake_logit = D(fake_img.detach())
g_loss = layers.mean(softplus(-fake_logit))
#print("g_loss",g_loss)
loss_G_list.append(g_loss.numpy())
# Update generator
g_loss.backward()
optim_G.minimize(g_loss)
# Output training stats
bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format(
ep, i + 1, 52000, loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
G.eval()
#fake_img = G(fix_z).detach().cpu()
fake_img = G(fix_z).numpy().squeeze()
log(f"fake_img.shape: {fake_img.shape}")
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'))
G.train()
# Save model
# print("type:",type(G.state_dict()).__name__)
# print("type:",type(D.state_dict()).__name__)
states = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
#dygraph.save_dygraph(state, os.path.join(opts.det, 'models', 'latest'))
save_checkpoint(states,
os.path.join(opts.det, 'models', 'latest.pp'))
# scheduler_D.step()
# scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)