def train(**kwargs):
vis = Visualizer('GAN') # 可视化
for epoch in range(num_epoch):
for i, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = img.cuda()
# train D
D_optimizer.zero_grad()
# 将 real_img 判断成 1
output = D(real_img)
D_loss_real = criterion(output, true_labels)
D_loss_real.backward()
# 将 fake_img 判断成 0
noises.data.copy_(torch.randn(batch_size, z_dimension, 1, 1))
fake_img = G(noises).detach() # draw fake pic
output = D(fake_img)
D_loss_fake = criterion(output, fake_labels)
D_loss_fake.backward()
D_optimizer.step()
D_loss = D_loss_real + D_loss_fake
errorD_meter.add(D_loss.item())
if i % 5 == 0: # train G every 5 batches
G_optimizer.zero_grad()
noises.data.copy_(torch.randn(batch_size, z_dimension, 1, 1))
fake_img = G(noises)
output = D(fake_img)
G_loss = criterion(output, true_labels)
G_loss.backward()
G_optimizer.step()
errorG_meter.add(G_loss.item())
if (epoch + 1) % 10 == 0: # save model every 10 epochs
if os.path.exists('/tmp/debugGAN'):
ipdb.set_trace()
fix_fake_imgs = G(fix_noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5,
win='fake image')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5,
win='real image')
vis.plot('errorD', errorD_meter.value()[0])
vis.plot('errorG', errorG_meter.value()[0])
torchvision.utils.save_image(fix_fake_imgs.data[:64],
'%s/%s.png' % ('imgs', epoch),
normalize=True,
range=(-1, 1))
torch.save(D.state_dict(), 'checkpoints/D_%s.pth' % epoch)
torch.save(G.state_dict(), 'checkpoints/G_%s.pth' % epoch)
errorD_meter.reset()
errorG_meter.reset()
def train(self, dataset):
if self.visual:
vis = Visualizer(self.env)
for epoch in range(self.epoch):
for ii, data in tqdm.tqdm(enumerate(dataset)):
real_img = data['I']
tri_img = data['T']
if self.com_loss:
bg_img = data['B'].to(self.device)
fg_img = data['F'].to(self.device)
# input to the G
input_img = t.tensor(np.append(real_img.numpy(), tri_img.numpy(), axis=1)).to(self.device)
# real_alpha
real_alpha = data['A'].to(self.device)
# vis.images(real_img.numpy()*0.5 + 0.5, win='input_real_img')
# vis.images(real_alpha.cpu().numpy()*0.5 + 0.5, win='real_alpha')
# vis.images(tri_img.numpy()*0.5 + 0.5, win='tri_map')
# train D
if ii % self.d_every == 0:
self.D_optimizer.zero_grad()
# real_img_d = input_img[:, 0:3, :, :]
tri_img_d = input_img[:, 3:4, :, :]
# 真正的alpha 交给判别器判断
if self.com_loss:
real_d = self.D(input_img)
else:
real_d = self.D(t.cat([real_alpha, tri_img_d], dim=1))
target_real_label = t.tensor(1.0)
target_real = target_real_label.expand_as(real_d).to(self.device)
loss_d_real = self.D_criterion(real_d, target_real)
#loss_d_real.backward()
# 生成器生成fake_alpha 交给判别器判断
fake_alpha = self.G(input_img)
if self.com_loss:
fake_img = fake_alpha*fg_img + (1 - fake_alpha) * bg_img
fake_d = self.D(t.cat([fake_img, tri_img_d], dim=1))
else:
fake_d = self.D(t.cat([fake_alpha, tri_img_d], dim=1))
target_fake_label = t.tensor(0.0)
target_fake = target_fake_label.expand_as(fake_d).to(self.device)
loss_d_fake = self.D_criterion(fake_d, target_fake)
loss_D = loss_d_real + loss_d_fake
loss_D.backward()
self.D_optimizer.step()
self.D_error_meter.add(loss_D.item())
# train G
if ii % self.g_every == 0:
self.G_optimizer.zero_grad()
real_img_g = input_img[:, 0:3, :, :]
tri_img_g = input_img[:, 3:4, :, :]
fake_alpha = self.G(input_img)
# fake_alpha 与 real_alpha的L1 loss
loss_g_alpha = self.G_criterion(fake_alpha, real_alpha)
loss_G = loss_g_alpha
self.Alpha_loss_meter.add(loss_g_alpha.item())
if self.com_loss:
fake_img = fake_alpha * fg_img + (1 - fake_alpha) * bg_img
loss_g_cmp = self.G_criterion(fake_img, real_img_g)
# 迷惑判别器
fake_d = self.D(t.cat([fake_img, tri_img_g], dim=1))
self.Com_loss_meter.add(loss_g_cmp.item())
loss_G = loss_G + loss_g_cmp
else:
fake_d = self.D(t.cat([fake_alpha, tri_img_g], dim=1))
target_fake = t.tensor(1.0).expand_as(fake_d).to(self.device)
loss_g_d = self.D_criterion(fake_d, target_fake)
self.Adv_loss_meter.add(loss_g_d.item())
loss_G = loss_G + loss_g_d
loss_G.backward()
self.G_optimizer.step()
self.G_error_meter.add(loss_G.item())
if self.visual and ii % 20 == 0:
vis.plot('errord', self.D_error_meter.value()[0])
#vis.plot('errorg', self.G_error_meter.value()[0])
vis.plot('errorg', np.array([self.Adv_loss_meter.value()[0], self.Alpha_loss_meter.value()[0],
self.Com_loss_meter.value()[0]]), legend=['adv_loss', 'alpha_loss',
'com_loss'])
vis.images(tri_img.numpy()*0.5 + 0.5, win='tri_map')
vis.images(real_img.cpu().numpy() * 0.5 + 0.5, win='relate_real_input')
vis.images(real_alpha.cpu().numpy() * 0.5 + 0.5, win='relate_real_alpha')
vis.images(fake_alpha.detach().cpu().numpy(), win='fake_alpha')
if self.com_loss:
vis.images(fake_img.detach().cpu().numpy()*0.5 + 0.5, win='fake_img')
self.G_error_meter.reset()
self.D_error_meter.reset()
self.Alpha_loss_meter.reset()
self.Com_loss_meter.reset()
self.Adv_loss_meter.reset()
if epoch % 5 == 0:
t.save(self.D.state_dict(), self.save_dir + '/netD' + '/netD_%s.pth' % epoch)
t.save(self.G.state_dict(), self.save_dir + '/netG' + '/netG_%s.pth' % epoch)
return
def train(self, dataset):
if self.visual:
vis = Visualizer(self.env)
for epoch in range(self.epoch):
for ii, data in tqdm.tqdm(enumerate(dataset)):
real_img = data['I']
tri_img = data['T'] #Trimap
if self.com_loss:
bg_img = data['B'].to(self.device) #Background image
fg_img = data['F'].to(self.device) #Foreground image
# input to the G, 4 Channel, Image and Trimap concatenated
input_img = t.tensor(
np.append(real_img.numpy(), tri_img.numpy(),
axis=1)).to(self.device)
# real_alpha
real_alpha = data['A'].to(self.device)
# vis.images(real_img.numpy()*0.5 + 0.5, win='input_real_img')
# vis.images(real_alpha.cpu().numpy()*0.5 + 0.5, win='real_alpha')
# vis.images(tri_img.numpy()*0.5 + 0.5, win='tri_map')
# train D
if ii % self.d_every == 0:
self.D_optimizer.zero_grad()
# real_img_d = input_img[:, 0:3, :, :]
tri_img_d = input_img[:, 3:4, :, :]
#alpha
if self.com_loss:
real_d = self.D(input_img)
else:
real_d = self.D(t.cat([real_alpha, tri_img_d], dim=1))
target_real_label = t.tensor(1.0) #1 for real
#The shape of real_d would be NxN
target_real = target_real_label.expand_as(real_d).to(
self.device)
loss_d_real = self.D_criterion(real_d, target_real)
#fake_alpha, is the predicted alpha by the generator
fake_alpha = self.G(input_img)
if self.com_loss:
#Constructing the fake Image
fake_img = fake_alpha * fg_img + (1 -
fake_alpha) * bg_img
fake_d = self.D(t.cat([fake_img, tri_img_d], dim=1))
else:
fake_d = self.D(t.cat([fake_alpha, tri_img_d], dim=1))
target_fake_label = t.tensor(0.0)
target_fake = target_fake_label.expand_as(fake_d).to(
self.device)
loss_d_fake = self.D_criterion(fake_d, target_fake)
loss_D = loss_d_real + loss_d_fake
#Backpropagation of the discriminator loss
loss_D.backward()
self.D_optimizer.step()
self.D_error_meter.add(loss_D.item())
# train G
if ii % self.g_every == 0:
#Initialize the Optimizer
self.G_optimizer.zero_grad()
real_img_g = input_img[:, 0:3, :, :]
tri_img_g = input_img[:, 3:4, :, :]
fake_alpha = self.G(input_img)
# fake_alpha is the output of the Generator
loss_g_alpha = self.G_criterion(fake_alpha, real_alpha)
#alpha_loss, difference between predicted alpha and the real alpha
loss_G = loss_g_alpha
self.Alpha_loss_meter.add(loss_g_alpha.item())
if self.com_loss:
fake_img = fake_alpha * fg_img + (1 -
fake_alpha) * bg_img
loss_g_cmp = self.G_criterion(
fake_img, real_img_g) #Composition Loss
fake_d = self.D(t.cat([fake_img, tri_img_g], dim=1))
self.Com_loss_meter.add(loss_g_cmp.item())
loss_G = loss_G + loss_g_cmp
else:
fake_d = self.D(t.cat([fake_alpha, tri_img_g], dim=1))
target_fake = t.tensor(1.0).expand_as(fake_d).to(
self.device)
#The target of Generator is to make the Discriminator ouptut 1
loss_g_d = self.D_criterion(fake_d, target_fake)
self.Adv_loss_meter.add(loss_g_d.item())
loss_G = loss_G + loss_g_d
loss_G.backward()
self.G_optimizer.step()
self.G_error_meter.add(loss_G.item())
if self.visual and ii % 20 == 0:
vis.plot('errord', self.D_error_meter.value()[0])
#vis.plot('errorg', self.G_error_meter.value()[0])
vis.plot('errorg',
np.array([
self.Adv_loss_meter.value()[0],
self.Alpha_loss_meter.value()[0],
self.Com_loss_meter.value()[0]
]),
legend=['adv_loss', 'alpha_loss', 'com_loss'])
vis.images(tri_img.numpy() * 0.5 + 0.5, win='tri_map')
vis.images(real_img.cpu().numpy() * 0.5 + 0.5,
win='relate_real_input')
vis.images(real_alpha.cpu().numpy() * 0.5 + 0.5,
win='relate_real_alpha')
vis.images(fake_alpha.detach().cpu().numpy(),
win='fake_alpha')
if self.com_loss:
vis.images(fake_img.detach().cpu().numpy() * 0.5 + 0.5,
win='fake_img')
self.G_error_meter.reset()
self.D_error_meter.reset()
self.Alpha_loss_meter.reset()
self.Com_loss_meter.reset()
self.Adv_loss_meter.reset()
if epoch % 5 == 0:
t.save(self.D.state_dict(),
self.save_dir + '/netD' + '/netD_%s.pth' % epoch)
t.save(self.G.state_dict(),
self.save_dir + '/netG' + '/netG_%s.pth' % epoch)
return
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device=t.device('cuda') if opt.gpu else t.device('cpu')
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env)
# 数据
transforms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True
)
# 网络
netg, netd = NetG(opt), NetD(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
print("use pretrained models.")
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
netd.to(device)
netg.to(device)
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = t.nn.BCELoss().to(device)
# 真图片label为1，假图片label为0
# noises为生成网络的输入
true_labels = t.ones(opt.batch_size).to(device)
fake_labels = t.zeros(opt.batch_size).to(device)
fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
epochs = range(opt.max_epoch)
for epoch in iter(epochs):
for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = img.to(device)
if ii % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
## 尽可能的把真图片判别为正确
output = netd(real_img)
error_d_real = criterion(output, true_labels)
error_d_real.backward()
## 尽可能把假图片判别为错误
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises).detach() # 根据噪声生成假图
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward()
optimizer_d.step()
error_d = error_d_fake + error_d_real
errord_meter.add(error_d.item())
if ii % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises)
output = netd(fake_img)
error_g = criterion(output, true_labels)
error_g.backward()
optimizer_g.step()
errorg_meter.add(error_g.item())
if opt.vis and ii % opt.plot_every == opt.plot_every - 1:
## 可视化
if os.path.exists(opt.debug_file):
ipdb.set_trace()
fix_fake_imgs = netg(fix_noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real')
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
if (epoch+1) % opt.save_every == 0:
# 保存模型、图片
# tv.utils.save_image(fix_fake_imgs.data[:64], '{0}_{1}.png'.format(opt.save_path, epoch), normalize=True, range=(-1, 1))
t.save(netd.state_dict(), './checkpoints/tiny_netd_%s.pth' % (epoch))
t.save(netg.state_dict(), './checkpoints/tiny_netg_%s.pth' % (epoch))
errord_meter.reset()
errorg_meter.reset()
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device = t.device('cuda') if opt.gpu else t.device('cpu')
# 可视化
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env)
# 数据
transforms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = tv.datasets.ImageFolder(opt.data_path, transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True)
netg, netd = NetG(opt), NetD(opt)
if opt.netd_path:
netd.load_state_dict(
t.load(opt.netd_path, map_location=t.device('cpu')))
if opt.netg_path:
netg.load_state_dict(
t.load(opt.netg_path, map_location=t.device('cpu')))
netd.to(device)
netg.to(device)
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(),
opt.lr1,
betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(),
opt.lr1,
betas=(opt.beta1, 0.999))
criterion = t.nn.BCELoss().to(device)
# 真图片label为1，假图片label为0
# noise为生成网络的输入
true_labels = t.ones(opt.batch_size).to(device)
fake_labels = t.zeros(opt.batch_size).to(device)
fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
errord = 0
errorg = 0
epochs = range(opt.max_epoch)
for epoch in epochs:
for ii, (img, _) in tqdm(enumerate(dataloader)):
real_img = img.to(device)
if (ii + 1) % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
# 把真图片判断为正确
output = netd(real_img)
error_d_real = criterion(output, true_labels)
error_d_real.backward()
# 把假图片(netg通过噪声生成的图片)判断为错误
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises).detach()
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward()
optimizer_d.step()
errord += (error_d_fake + error_d_real).item()
if (ii + 1) % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises)
output = netd(fake_img)
error_g = criterion(output, true_labels)
error_g.backward()
optimizer_g.step()
errorg += error_g.item()
if opt.vis and (ii + 1) % opt.plot_every == 0:
fix_fake_imgs = netg(fix_noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 +
0.5,
win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5,
win='real')
vis.plot('errord', errord / (opt.plot_every))
vis.plot('errorg', errorg / (opt.plot_every))
errord = 0
errorg = 0
if (epoch + 1) % opt.save_every == 0:
tv.utils.save_image(fix_fake_imgs.data[:64],
'%s%s.png' % (opt.save_path, epoch),
normalize=True,
range=(-1, 1))
t.save(netd.state_dict(),
'checkpoints/netd_%s.pth' % (epoch + 1))
t.save(netg.state_dict(),
'checkpoints/netg_%s.pth' % (epoch + 1))
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device=t.device('cuda') if opt.gpu else t.device('cpu')
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env)
# 数据
transforms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True
)
# 网络
netg, netd = NetG(opt), NetD(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
netd.to(device)
netg.to(device)
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = t.nn.BCELoss().to(device)
# 真图片label为1，假图片label为0
# noises为生成网络的输入
true_labels = t.ones(opt.batch_size).to(device)
fake_labels = t.zeros(opt.batch_size).to(device)
fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
epochs = range(opt.max_epoch)
for epoch in iter(epochs):
for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = img.to(device)
if ii % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
## 尽可能的把真图片判别为正确
output = netd(real_img)
error_d_real = criterion(output, true_labels)
error_d_real.backward()
## 尽可能把假图片判别为错误
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises).detach() # 根据噪声生成假图
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward()
optimizer_d.step()
error_d = error_d_fake + error_d_real
errord_meter.add(error_d.item())
if ii % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises)
output = netd(fake_img)
error_g = criterion(output, true_labels)
error_g.backward()
optimizer_g.step()
errorg_meter.add(error_g.item())
if opt.vis and ii % opt.plot_every == opt.plot_every - 1:
## 可视化
if os.path.exists(opt.debug_file):
ipdb.set_trace()
fix_fake_imgs = netg(fix_noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real')
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
if (epoch+1) % opt.save_every == 0:
# 保存模型、图片
tv.utils.save_image(fix_fake_imgs.data[:64], '%s/%s.png' % (opt.save_path, epoch), normalize=True,
range=(-1, 1))
t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch)
t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch)
errord_meter.reset()
errorg_meter.reset()
def train(**kwargs):
for k_,v_ in kwargs.items():
setattr(opt,k_,v_)
if opt.vis is True:
from visualize import Visualizer
vis = Visualizer(opt.env)
transforms = tv.transforms.Compose([
tv.transforms.Scale(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
dataset = tv.datasets.ImageFolder(opt.data_path,transform=transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size= opt.batch_size,
shuffle = True,
num_workers = opt.num_workers,
drop_last = True)
netg,netd = NetG(opt),NetD(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
netd.load_state_dict(t.load(opt.netd_path,map_location = map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path,map_location= map_location))
optimizer_g = t.optim.Adam(netg.parameters(),opt.lr1,betas=(opt.beta1,0.999))
optimizer_d = t.optim.Adam(netd.parameters(),opt.lr2,betas=(opt.beta1,0.999))
criterion = t.nn.BCELoss()
true_labels = Variable(t.ones(opt.batch_size))
fake_labels = Variable(t.zeros(opt.batch_size))
fix_noises = Variable(t.randn(opt.batch_size,opt.nz,1,1))
noises = Variable(t.randn(opt.batch_size,opt.nz,1,1))
if opt.use_gpu:
netd.cuda()
netg.cuda()
criterion.cuda()
true_labels,fake_labels = true_labels.cuda(),fake_labels.cuda()
fix_noises,noises = fix_noises.cuda(),noises.cuda()
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
epoches =range(opt.max_epoch)
for epoch in iter(epoches):
for ii, (img,_) in tqdm(enumerate(dataloader)):
real_img = Variable(img)
if opt.use_gpu:
real_img = real_img.cuda()
if (ii +1 )% opt.d_every == 0:
optimizer_d.zero_grad()
output = netd(real_img)
error_d_real = criterion(output,true_labels)
error_d_real.backward()
noises.data.copy_(t.randn(opt.batch_size,opt.nz,1,1))
fake_img = netg(noises).detach()
fake_output = netd(fake_img)
error_d_fake = criterion(fake_output,fake_labels)
error_d_fake.backward()
optimizer_d.step()
error_d =error_d_real + error_d_fake
errord_meter.add(error_d.data)
if (ii + 1)%opt.g_every == 0:
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size,opt.nz,1,1))
fake_img = netg(noises)
fake_output = netd(fake_img)
error_g = criterion(fake_output,true_labels)
error_g.backward()
optimizer_g.step()
errorg_meter.add(error_g.data)
if opt.vis and ii%opt.plot_every == opt.plot_every -1 :
#visualize
if os.path.exists(opt.debug_file):
ipdb.set_trace()
fix_fake_img = netg(fix_noises)
vis.images(fix_fake_img.data.cpu().numpy()[:16]*0.5 + 0.5 ,win='fixfake')
vis.images(real_img.data.cpu().numpy()[:16]*0.5+0.5,win='real')
vis.plot('errord',errord_meter.value()[0])
vis.plot('errorg',errorg_meter.value()[0])
if (epoch+1) % opt.save_every ==0:
tv.utils.save_image(fix_fake_img.data[:64],'%s/%s.png'%(opt.save_path,epoch),normalize=True,range=(-1,1))
t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch)
t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch)
errord_meter.reset()
errorg_meter.reset()
def train(self, dataset):
vis = Visualizer('alphaGAN_train')
for epoch in range(self.epoch):
for ii, data in tqdm.tqdm(enumerate(dataset)):
real_img = data['I'].to(self.device)
tri_img = data['T']
bg_img = data['B'].to(self.device)
fg_img = data['F'].to(self.device)
# input to the G
input_img = t.tensor(
np.append(real_img.cpu().numpy(), tri_img.numpy(),
axis=1)).to(self.device)
# real_alpha
real_alpha = data['A'].to(self.device)
# vis.images(real_img.numpy()*0.5 + 0.5, win='input_real_img')
# vis.images(real_alpha.cpu().numpy()*0.5 + 0.5, win='real_alpha')
# vis.images(tri_img.numpy()*0.5 + 0.5, win='tri_map')
if ii % 5 == 0:
self.D_optimizer.zero_grad()
# 真正的alpha 交给判别器判断
if self.com_loss:
real_d = self.D(real_img)
else:
real_d = self.D(real_alpha)
target_real_label = t.tensor(1.0)
target_real = target_real_label.expand_as(real_d).to(
self.device)
loss_d_real = self.D_criterion(real_d, target_real)
#loss_d_real.backward()
# 生成器生成fake_alpha 交给判别器判断
fake_alpha = self.G(input_img)
if self.com_loss:
fake_img = fake_alpha * fg_img + (1 -
fake_alpha) * bg_img
fake_d = self.D(fake_img)
else:
fake_d = self.D(fake_alpha)
target_fake_label = t.tensor(0.0)
target_fake = target_fake_label.expand_as(fake_d).to(
self.device)
loss_d_fake = self.D_criterion(fake_d, target_fake)
loss_D = (loss_d_real + loss_d_fake) * 0.5
loss_D.backward()
self.D_optimizer.step()
self.D_error_meter.add(loss_D.item())
if ii % 1 == 0:
self.G_optimizer.zero_grad()
fake_alpha = self.G(input_img)
# fake_alpha 与 real_alpha的L1 loss
loss_g_alpha = self.G_criterion(fake_alpha, real_alpha)
loss_G = 0.8 * loss_g_alpha
if self.com_loss:
fake_img = fake_alpha * fg_img + (1 -
fake_alpha) * bg_img
loss_g_cmp = self.G_criterion(fake_img, real_img)
# 迷惑判别器
fake_d = self.D(fake_img)
loss_G = loss_G + 0.2 * loss_g_cmp
else:
fake_d = self.D(fake_alpha)
target_fake = t.tensor(1.0).expand_as(fake_d).to(
self.device)
loss_g_d = self.D_criterion(fake_d, target_fake)
loss_G = loss_G + 0.2 * loss_g_d
loss_G.backward()
self.G_optimizer.step()
self.G_error_meter.add(loss_G.item())
if ii % 20 == 0:
vis.plot('errord', self.D_error_meter.value()[0])
vis.plot('errorg', self.G_error_meter.value()[0])
vis.images(tri_img.numpy() * 0.5 + 0.5, win='tri_map')
vis.images(real_img.cpu().numpy() * 0.5 + 0.5,
win='relate_real_input')
vis.images(real_alpha.cpu().numpy() * 0.5 + 0.5,
win='relate_real_alpha')
vis.images(fake_alpha.detach().cpu().numpy(),
win='fake_alpha')
if self.com_loss:
vis.images(fake_img.detach().cpu().numpy() * 0.5 + 0.5,
win='fake_img')
self.G_error_meter.reset()
self.D_error_meter.reset()
if epoch % 5 == 0:
t.save(
self.D.state_dict(), '/home/zzl/model/alphaGAN/netD/' +
self.save_dir + '/netD_%s.pth' % epoch)
t.save(
self.G.state_dict(), '/home/zzl/model/alphaGAN/netG/' +
self.save_dir + '/netG_%s.pth' % epoch)
return
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device = t.device('cuda') if opt.gpu else t.device('cpu')
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env, server="172.16.6.194")
# 数据
transforms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = mydataset.AADBDataset(opt.data_path, opt.lable_path, transforms=transforms)
# dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True
)
# 网络
netg, netd = NetmyG(opt), NetD(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
netd.to(device)
netg.to(device)
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = t.nn.BCELoss().to(device)
# huber损失
huber = t.nn.SmoothL1Loss().to(device)
# 真图片label为1，假图片label为0
# noises为生成网络的输入
# TODO
true_labels = t.ones([opt.batch_size, opt.weidu]).to(device)
fake_labels = t.zeros([opt.batch_size, opt.weidu]).to(device)
# noises 为12维
fix_noises = t.randn(opt.batch_size, opt.nz, 1, opt.weidu).to(device)
noises = t.randn(opt.batch_size, opt.nz, 1, opt.weidu).to(device)
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
errorgs_meter = AverageValueMeter()
epochs = range(opt.max_epoch)
for epoch in iter(epochs):
for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = img.to(device)
if ii % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
## 尽可能的把真图片判别为正确
output = netd(real_img)
error_d_real = criterion(output, true_labels)
error_d_real.backward(retain_graph=True)
## 尽可能把假图片判别为错误
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, opt.weidu))
fake_img = netg(noises).detach() # 根据噪声生成假图
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward(retain_graph=True)
optimizer_d.step()
error_d = error_d_fake + error_d_real
errord_meter.add(error_d.item())
if ii % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, opt.weidu))
fake_img = netg(noises)
output = netd(fake_img)
error_g = criterion(output, true_labels)
print(output.shape, true_labels.shape)
error_g.backward(retain_graph=True)
bb = noises.view(opt.batch_size,-1)
print (bb.shape)
error_s = huber(output, bb)
error_s.backward(retain_graph=True)
error_gs = error_s + error_g
optimizer_g.step()
errorg_meter.add(error_g.item())
errorgs_meter.add(error_gs.item())
if opt.vis and ii % opt.plot_every == opt.plot_every - 1:
## 可视化
if os.path.exists(opt.debug_file):
ipdb.set_trace()
fix_fake_imgs = netg(fix_noises).detach()
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real')
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
vis.plot('errorgs', errorgs_meter.value()[0])
if (epoch + 1) % opt.save_every == 0:
# 保存模型、图片
localtime = time.asctime(time.localtime(time.time()))
print("要保存了")
t.save(netd.state_dict(), 'chenck20190529/netd_%s.pth' % epoch)
t.save(netg.state_dict(), 'chenck20190529/netg_%s.pth' % epoch)
errord_meter.reset()
errorg_meter.reset()
errorgs_meter.reset()
def train(**kwargs):
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
device = t.device('cuda') if opt.gpu else t.device('cpu')
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env)
# 数据
transforms = tv.transforms.Compose([
tv.transforms.Resize(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms)
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True
)
# 网络
netg, netd = NetG(opt), NetD(opt)
map_location = lambda storage, loc: storage
if opt.netd_path:
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
# 把网络放入GPU中，如果有的话
netd.to(device)
netg.to(device)
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
# binart cross entropy
criterion = t.nn.BCELoss().to(device)
# 真图片label为1，假图片label为0，label是batchsize长度的数组
# noises为生成网络的输入
true_labels = t.ones(opt.batch_size).to(device)
fake_labels = t.zeros(opt.batch_size).to(device)
fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
epochs = range(opt.max_epoch)
for epoch in iter(epochs):
for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = img.to(device)
if ii % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
## 尽可能的把真图片判别为正确
output = netd(real_img)
# 将real image放入，然后loss和true label求取
error_d_real = criterion(output, true_labels)
error_d_real.backward()
## 尽可能把假图片判别为错误，注意这里的加图片是noise从generator中拿到的！！！
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises).detach() # 根据噪声生成假图
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward()
# 上面两个都反向传播误差之后，才进行更新（optimizer自动获取parameter里面的grad然后更新）
# ！！！非常重要，这里只有判别器的parameter更新了，generator的权重不会更新！
optimizer_d.step()
error_d = error_d_fake + error_d_real
errord_meter.add(error_d.item())
if ii % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises)
output = netd(fake_img)
# fake image经过generater之后变成fake image，之后这个image会努力使其被判定为true
# 但是需要注意，这里优化的参数只有生成器的，没有判别器的
error_g = criterion(output, true_labels)
error_g.backward()
optimizer_g.step()
errorg_meter.add(error_g.item())
if opt.vis and ii % opt.plot_every == opt.plot_every - 1:
## 可视化
if os.path.exists(opt.debug_file):
ipdb.set_trace()
fix_fake_imgs = netg(fix_noises)
vis.images(fix_fake_imgs.detach().cpu().numpy()[:64] * 0.5 + 0.5, win='fixfake')
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5, win='real')
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
if (epoch + 1) % opt.save_every == 0:
# 保存模型、图片
tv.utils.save_image(fix_fake_imgs.data[:64], '%s/%s.png' % (opt.save_path, epoch), normalize=True,
range=(-1, 1))
t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch)
t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch)
errord_meter.reset()
errorg_meter.reset()
def train(**kwargs):
# 读取参数赋值
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
# 可视化
if opt.vis:
from visualize import Visualizer
vis = Visualizer(opt.env)
# 对图片进行操作
transforms = tv.transforms.Compose([
tv.transforms.Scale(opt.image_size),
tv.transforms.CenterCrop(opt.image_size),
tv.transforms.ToTensor(),
# 均值方差
tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# ImageFolder 使用pytorch原生的方法读取图片，并进行操作 封装数据集
dataset = tv.datasets.ImageFolder(opt.data_path, transform=transforms)
#数据加载器
dataloader = t.utils.data.DataLoader(dataset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True)
# 定义网络
netg, netd = NetG(opt), NetD(opt)
# 把map内容加载到CPU中
map_location = lambda storage, loc: storage
# 将预训练的模型都先加载到cpu上
if opt.netd_path:
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
# 定义优化器和损失
optimizer_g = t.optim.Adam(netg.parameters(),
opt.lr1,
betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(),
opt.lr2,
betas=(opt.beta1, 0.999))
# BinaryCrossEntropy二分类交叉熵，常用于二分类问题，当然也可以用于多分类问题
criterion = t.nn.BCELoss()
# 真图片label为1，假图片label为0
# noises为生成网络的输入
true_labels = Variable(t.ones(opt.batch_size))
fake_labels = Variable(t.zeros(opt.batch_size))
# fix_noises是固定值，用来查看每个epoch的变化效果
fix_noises = Variable(t.randn(opt.batch_size, opt.nz, 1, 1))
noises = Variable(t.randn(opt.batch_size, opt.nz, 1, 1))
# AverageValueMeter统计任意添加的变量的方差和均值 可视化的仪表盘
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
if opt.gpu:
# 网络转移到GPU
netd.cuda()
netg.cuda()
# 损失函数转移到GPU
criterion.cuda()
# 标签转移到GPU
true_labels, fake_labels = true_labels.cuda(), fake_labels.cuda()
# 输入噪声转移到GPU
fix_noises, noises = fix_noises.cuda(), noises.cuda()
epochs = range(opt.max_epoch)
for epoch in iter(epochs):
for ii, (img, _) in tqdm.tqdm(enumerate(dataloader)):
real_img = Variable(img)
if opt.gpu:
real_img = real_img.cuda()
# 每d_every个batch训练判别器
if ii % opt.d_every == 0:
# 训练判别器
optimizer_d.zero_grad()
## 尽可能的把真图片判别为正确
#一个batchd的真照片判定为1 并反向传播
output = netd(real_img)
error_d_real = criterion(output, true_labels)
#反向传播
error_d_real.backward()
## 尽可能把假图片判别为错误
# 一个batchd的假照片判定为0 并反向传播
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
fake_img = netg(noises).detach() # 根据噪声生成假图
output = netd(fake_img)
error_d_fake = criterion(output, fake_labels)
error_d_fake.backward()
#更新可学习参数
optimizer_d.step()
# 总误差=识别真实图片误差+假图片误差
error_d = error_d_fake + error_d_real
# 将总误差加入仪表板用于可视化显示
errord_meter.add(error_d.data[0])
# 每g_every个batch训练生成器
if ii % opt.g_every == 0:
# 训练生成器
optimizer_g.zero_grad()
noises.data.copy_(t.randn(opt.batch_size, opt.nz, 1, 1))
# 生成器：噪声生成假图片
fake_img = netg(noises)
# 判别器：假图片判别份数
output = netd(fake_img)
# 尽量让假图片的份数与真标签接近,让判别器分不出来
error_g = criterion(output, true_labels)
error_g.backward()
# 更新参数
optimizer_g.step()
# 将误差加入仪表板用于可视化显示
errorg_meter.add(error_g.data[0])
if opt.vis and ii % opt.plot_every == opt.plot_every - 1:
## 可视化
# 进入debug模式
if os.path.exists(opt.debug_file):
ipdb.set_trace()
# 固定噪声生成假图片
fix_fake_imgs = netg(fix_noises)
# 可视化 固定噪声产生的假图片
vis.images(fix_fake_imgs.data.cpu().numpy()[:64] * 0.5 + 0.5,
win='fixfake')
# 可视化一张真图片。作为对比
vis.images(real_img.data.cpu().numpy()[:64] * 0.5 + 0.5,
win='real')
# 可视化仪表盘 判别器误差 生成器误差
vis.plot('errord', errord_meter.value()[0])
vis.plot('errorg', errorg_meter.value()[0])
# 每decay_every个epoch之后保存一次模型
if epoch % opt.decay_every == 0:
# 保存模型、图片
tv.utils.save_image(fix_fake_imgs.data[:64],
'%s/%s.png' % (opt.save_path, epoch),
normalize=True,
range=(-1, 1))
# 保存判别器 生成器
t.save(netd.state_dict(), 'checkpoints/netd_%s.pth' % epoch)
t.save(netg.state_dict(), 'checkpoints/netg_%s.pth' % epoch)
# 清空误差仪表盘
errord_meter.reset()
errorg_meter.reset()
# 重置优化器参数为刚开始的参数
optimizer_g = t.optim.Adam(netg.parameters(),
opt.lr1,
betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(),
opt.lr2,
betas=(opt.beta1, 0.999))
