def main():
# load training data
trainset = Dataset('./data/brilliant_blue')
trainloader = torch.utils.data.DataLoader(
trainset, batch_size=batch_size, shuffle=True
)
# init netD and netG
netD = Discriminator().to(device)
netD.apply(weights_init)
netG = Generator(nz).to(device)
netG.apply(weights_init)
criterion = nn.BCELoss()
# used for visualzing training process
fixed_noise = torch.randn(16, nz, 1, device=device)
real_label = 1.
fake_label = 0.
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))
for epoch in range(epoch_num):
for step, (data, _) in enumerate(trainloader):
real_cpu = data.to(device)
b_size = real_cpu.size(0)
# train netD
label = torch.full((b_size,), real_label,
dtype=torch.float, device=device)
netD.zero_grad()
output = netD(real_cpu).view(-1)
errD_real = criterion(output, label)
errD_real.backward()
D_x = output.mean().item()
# train netG
noise = torch.randn(b_size, nz, 1, device=device)
fake = netG(noise)
label.fill_(fake_label)
output = netD(fake.detach()).view(-1)
errD_fake = criterion(output, label)
errD_fake.backward()
D_G_z1 = output.mean().item()
errD = errD_real + errD_fake
optimizerD.step()
netG.zero_grad()
label.fill_(real_label)
output = netD(fake).view(-1)
errG = criterion(output, label)
errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
% (epoch, epoch_num, step, len(trainloader),
errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
# save training process
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
f, a = plt.subplots(4, 4, figsize=(8, 8))
for i in range(4):
for j in range(4):
a[i][j].plot(fake[i * 4 + j].view(-1))
a[i][j].set_xticks(())
a[i][j].set_yticks(())
plt.savefig('./img/dcgan_epoch_%d.png' % epoch)
plt.close()
# save models
torch.save(netG, './nets/dcgan_netG.pkl')
torch.save(netD, './nets/dcgan_netD.pkl')
nearest_k = 5
print('dataloader size',len(dataloader.dataset))
for epoch in range(params['nepochs']):
for i, data in enumerate(dataloader, 0):
# Transfer data tensor to GPU/CPU (device)]\
real_data = data[0].to(device)
print('data type and len',type(real_data),len(real_data))
# Get batch size. Can be different from params['nbsize'] for last batch in epoch.
b_size = real_data.size(0)
# Make accumalated gradients of the discriminator zero.
netD.zero_grad()
# Create labels for the real data. (label=1)
label = torch.full((b_size, ), real_label, device=device)
#print("real",len(real_data))
output = netD(real_data).view(-1)
#print("output",len(output))
errD_real = criterion(output, label)
# Calculate gradients for backpropagation.
errD_real.backward()
D_x = output.mean().item()
# Sample random data from a unit normal distribution.
noise = torch.randn(b_size, params['nz'], 1, 1, device=device)
# Generate fake data (images).
fake_data = netG(noise)
def main():
# load training data
trainset = Dataset('./data/brilliant_blue')
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True)
# init netD and netG
netD = Discriminator().to(device)
netD.apply(weights_init)
netG = Generator(nz).to(device)
netG.apply(weights_init)
# used for visualizing training process
fixed_noise = torch.randn(16, nz, 1, device=device)
# optimizers
optimizerD = optim.RMSprop(netD.parameters(), lr=lr)
optimizerG = optim.RMSprop(netG.parameters(), lr=lr)
for epoch in range(epoch_num):
for step, (data, _) in enumerate(trainloader):
# training netD
real_cpu = data.to(device)
b_size = real_cpu.size(0)
netD.zero_grad()
noise = torch.randn(b_size, nz, 1, device=device)
fake = netG(noise)
loss_D = -torch.mean(netD(real_cpu)) + torch.mean(netD(fake))
loss_D.backward()
optimizerD.step()
for p in netD.parameters():
p.data.clamp_(-clip_value, clip_value)
if step % n_critic == 0:
# training netG
noise = torch.randn(b_size, nz, 1, device=device)
netG.zero_grad()
fake = netG(noise)
loss_G = -torch.mean(netD(fake))
netD.zero_grad()
netG.zero_grad()
loss_G.backward()
optimizerG.step()
if step % 5 == 0:
print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f' %
(epoch, epoch_num, step, len(trainloader), loss_D.item(),
loss_G.item()))
# save training process
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
f, a = plt.subplots(4, 4, figsize=(8, 8))
for i in range(4):
for j in range(4):
a[i][j].plot(fake[i * 4 + j].view(-1))
a[i][j].set_xticks(())
a[i][j].set_yticks(())
plt.savefig('./img/wgan_epoch_%d.png' % epoch)
plt.close()
# save model
torch.save(netG, './nets/wgan_netG.pkl')
torch.save(netD, './nets/wgan_netD.pkl')
for epoch in range(num_epochs):
for batch_idx, (real, _) in enumerate(loader):
real = real.to(device)
noise = torch.randn((batch_size, z_dim, 1, 1)).to(device)
fake = gen(noise)
# Training Discriminator max log(D(x)) + log(1 - D(G(z)))
disc_real = disc(real).view(-1)
loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))
disc_fake = disc(fake).view(-1)
loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
loss_disc = (loss_disc_real + loss_disc_fake) / 2
disc.zero_grad()
loss_disc.backward(retain_graph=True)
opt_disc.step()
# Training Generator min log(1 - D(G(z))) -> max log(D(G(z)))
output = disc(fake).view(-1)
loss_gen = criterion(output, torch.ones_like(output))
gen.zero_grad()
loss_gen.backward()
opt_gen.step()
if batch_idx % 200 == 0:
print(
"Epoch {}, Loss Discriminator: {}, Loss Generator: {}".format(
epoch, loss_disc, loss_gen))
gamma=0.1)
# step 5: iteration
img_list = []
G_losses = []
D_losses = []
iters = 0
for epoch in range(num_epochs):
for i, data in enumerate(train_loader):
############################
# (1) Update D network
###########################
net_d.zero_grad()
# create training data
real_img = data.to(device)
b_size = real_img.size(0)
# 根据 (b_size,) 构造矩阵，使用 real_idx填充
real_label = torch.full((b_size, ), real_idx, device=device)
noise = torch.randn(b_size, nz, 1, 1, device=device)
fake_img = net_g(noise)
fake_label = torch.full((b_size, ), fake_idx, device=device)
# train D with real img
out_d_real = net_d(real_img)
loss_d_real = criterion(out_d_real.view(-1), real_label)
def main():
# load training data
trainset = Dataset('./data/brilliant_blue')
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True)
# init netD and netG
netD = Discriminator().to(device)
netD.apply(weights_init)
netG = Generator(nz).to(device)
netG.apply(weights_init)
# used for visualizing training process
fixed_noise = torch.randn(16, nz, 1, device=device)
# optimizers
# optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, beta2))
# optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, beta2))
optimizerD = optim.RMSprop(netD.parameters(), lr=lr)
optimizerG = optim.RMSprop(netG.parameters(), lr=lr)
for epoch in range(epoch_num):
for step, (data, _) in enumerate(trainloader):
# training netD
real_cpu = data.to(device)
b_size = real_cpu.size(0)
netD.zero_grad()
noise = torch.randn(b_size, nz, 1, device=device)
fake = netG(noise)
# gradient penalty
eps = torch.Tensor(b_size, 1, 1).uniform_(0, 1)
x_p = eps * data + (1 - eps) * fake
grad = autograd.grad(netD(x_p).mean(),
x_p,
create_graph=True,
retain_graph=True)[0].view(b_size, -1)
grad_norm = torch.norm(grad, 2, 1)
grad_penalty = p_coeff * torch.pow(grad_norm - 1, 2)
loss_D = torch.mean(netD(fake) - netD(real_cpu))
loss_D.backward()
optimizerD.step()
for p in netD.parameters():
p.data.clamp_(-0.01, 0.01)
if step % n_critic == 0:
# training netG
noise = torch.randn(b_size, nz, 1, device=device)
netG.zero_grad()
fake = netG(noise)
loss_G = -torch.mean(netD(fake))
netD.zero_grad()
netG.zero_grad()
loss_G.backward()
optimizerG.step()
if step % 5 == 0:
print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f' %
(epoch, epoch_num, step, len(trainloader), loss_D.item(),
loss_G.item()))
# save training process
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
f, a = plt.subplots(4, 4, figsize=(8, 8))
for i in range(4):
for j in range(4):
a[i][j].plot(fake[i * 4 + j].view(-1))
a[i][j].set_xticks(())
a[i][j].set_yticks(())
plt.savefig('./img/wgan_gp_epoch_%d.png' % epoch)
plt.close()
# save model
torch.save(netG, './nets/wgan_gp_netG.pkl')
torch.save(netD, './nets/wgan_gp_netD.pkl')
def train():
os.makedirs('log', exist_ok=True)
ds = datasets.ImageFolder(root=data_root,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
dataloader = DataLoader(ds, batch_size=batch_size, shuffle=True)
net_g = Generator(n_latent_vector, n_g_filters).to(device)
net_g.apply(weight_init)
net_d = Discriminator(n_d_filters).to(device)
net_d.apply(weight_init)
if os.path.exists(model_save_path):
all_state_dict = torch.load(model_save_path)
net_d.load_state_dict(all_state_dict['d_state_dict'])
net_g.load_state_dict(all_state_dict['g_state_dict'])
print('model restored from {}'.format(model_save_path))
criterion = nn.BCELoss()
fixed_noise = torch.randn(1, n_latent_vector, 1, 1, device=device)
real_label = 1
fake_label = 0
optimizer_d = optim.Adam(net_d.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_g = optim.Adam(net_g.parameters(), lr=lr, betas=(0.5, 0.999))
print('start training...')
try:
for epoch in range(epochs):
for i, data in enumerate(dataloader, 0):
# update Discrinimator, maximize d loss
net_d.zero_grad()
real_cpu = data[0].to(device)
b_size = real_cpu.size(0)
label = torch.full((b_size,), real_label, device=device)
output = net_d(real_cpu).view(-1)
err_d_real = criterion(output, label)
err_d_real.backward()
d_x = output.mean().item()
# train with fake batch
noise = torch.randn(b_size, n_latent_vector, 1, 1, device=device)
fake = net_g(noise)
label.fill_(fake_label)
output = net_d(fake.detach()).view(-1)
err_d_fake = criterion(output, label)
err_d_fake.backward()
d_g_z1 = output.mean().item()
err_d = err_d_real + err_d_fake
optimizer_d.step()
# update Generator
net_g.zero_grad()
label.fill_(real_label)
output = net_d(fake).view(-1)
err_g = criterion(output, label)
err_g.backward()
d_g_z2 = output.mean().item()
optimizer_d.step()
if i % 50 == 0:
print(f'Epoch: {epoch}, loss_d: {err_d.item()}, loss_g: {err_g.item()}')
if epoch % 2 == 0 and epoch != 0:
with torch.no_grad():
fake = net_g(fixed_noise).detach().cpu().numpy()
print(fake.shape)
fake = np.transpose(np.squeeze(fake, axis=0), (1, 2, 0))
print(fake.shape)
cv2.imwrite('log/{}_fake.png'.format(epoch), fake)
print('record a fake image to local.')
except KeyboardInterrupt:
print('interrupted, try saving the model')
all_state_dict = {
'd_state_dict': net_d.state_dict(),
'g_state_dict': net_g.state_dict(),
}
torch.save(all_state_dict, model_save_path)
print('model saved...')
