import sys
sys.path.append("/home/hxj/anaconda3/lib/python3.6/site-packages")
from torchvision.datasets import ImageFolder
from PIL import Image
import torch
import os
import random
c_dim=5 # dimension of domain labels (1st dataset)
c2_dim=8 # dimension of domain labels (2nd dataset)
celeba_crop_size=178 # crop size for the CelebA dataset
rafd_crop_size=256 #crop size for the RaFD dataset
image_size=128 #image resolution
g_conv_dim=64 # number of conv filters in the first layer of G
d_conv_dim=64 # number of conv filters in the first layer of D
g_repeat_num = 6 #number of residual blocks in G
d_repeat_num=6 #number of strided conv layers in D
lambda_cls=1 #weight for domain classification loss
lambda_rec=10 # weight for reconstruction loss
lambda_gp=10 #'weight for gradient penalty
# Training configuration.
dataset='CelebA' # choices=['CelebA', 'RaFD', 'Both'])
batch_size=16 # 'mini-batch size
num_iters=200000 #number of total iterations for training D
num_iters_decay=100000 #number of iterations for decaying lr
g_lr=0.0001 #learning rate for G
d_lr=0.0001 #learning rate for D
n_critic=5 #number of D updates per each G update
beta1=0.5 #beta1 for Adam optimizer
beta2=0.999 #beta2 for Adam optimizer
resume_iters=None #resume training from this step
selected_attrs=['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Male', 'Young']
#selected attributes for the CelebA dataset'
# Test configuration.
test_iters=200000 #test model from this step
# Miscellaneous.
num_workers=1
mode='test' # choices=['train', 'test'])
use_tensorboard=True
# Directories.
celeba_image_dir='../data/CelebA_nocrop/images/' if mode == 'train' else '../test/test/'
attr_path='../data/list_attr_celeba.txt' if mode == 'train' else '../test/test_celeba.txt'
rafd_image_dir='../RaFD/train'
log_dir='../test/logs'
model_save_dir='../stargan/models'
sample_dir='../test/samples'
result_dir='../test/result'
# Step size.
log_step=10
sample_step=1000
model_save_step=10000
lr_update_step=1000
import tensorflow as tf
#这是加载TensorBord
class Logger(object):
"""Tensorboard logger."""
def __init__(self, log_dir):
"""Initialize summary writer."""
self.writer = tf.summary.FileWriter(log_dir)
def scalar_summary(self, tag, value, step):
"""Add scalar summary."""
summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
self.writer.add_summary(summary, step)
from torch.utils import data
from torchvision import transforms as T
class MulPIE(data.Dataset):
"""Dataset class for the MultiPIE dataset."""
def __init__(self, image_dir, transform, mode):
"""Initialize and preprocess the CelebA dataset."""
self.image_dir = image_dir
self.transform = transform
self.mode = mode
self.train_dataset = []
self.test_dataset = []
self.preprocess()
if mode == 'train':
self.num_images = len(self.train_dataset)
else:
self.num_images = len(self.test_dataset)
def preprocess(self):
dataset1 = ImageFolder(self.image_dir, self.transform)
for names in dataset1.imgs:
lable=[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
name = names[0]
num = int(name[-6:-4])
lable[num]=1
label_trg = [0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0]
self.train_dataset.append([name, lable, label_trg])
self.test_dataset.append([name, lable, label_trg])
#该方法是继承torch里面的utils文件夹里面data文件夹里面的Dataset类
def __getitem__(self, index):
"""Return one image and its corresponding attribute label."""
dataset = self.train_dataset if self.mode == 'train' else self.test_dataset
filename, label, label_trg= dataset[index]
#filename, label = dataset[index]
image = Image.open(os.path.join(self.image_dir, filename))
return self.transform(image), torch.FloatTensor(label), torch.FloatTensor(label_trg)
#return self.transform(image), torch.FloatTensor(label)
def __len__(self):
"""Return the number of images."""
return self.num_images
def get_loader(image_dir, image_size=128, batch_size=16, mode='train', num_workers=1):
"""Build and return a data loader."""
transform = []
if mode == 'train':
transform.append(T.RandomHorizontalFlip())
#transform1.append(T.CenterCrop(178)) 以后会用到裁剪图像
#to run only once
transform.append(T.Resize(image_size))
transform.append(T.ToTensor())
transform.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
transform = T.Compose(transform)
dataset = MulPIE(image_dir, transform, mode)
data_loader = data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=(mode=='train'),
num_workers=num_workers)
return data_loader
image_dir='../RaFD/'
mul_loader = get_loader(image_dir,image_size,batch_size, mode, num_workers)
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class ResidualBlock(nn.Module):
"""Residual Block with instance normalization."""
def __init__(self, dim_in, dim_out):
super(ResidualBlock, self).__init__()
self.main = nn.Sequential(
nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True),
nn.ReLU(inplace=True),
nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True))
def forward(self, x):
return x + self.main(x)
class Generator(nn.Module):
"""Generator network."""
def __init__(self, conv_dim=64, c_dim=5, repeat_num=6):
super(Generator, self).__init__()
layers = []
# 第一个卷积层,输入为图像和label的串联,3表示图像为3通道,c_dim为label的维度,
layers.append(nn.Conv2d(3+c_dim, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
layers.append(nn.InstanceNorm2d(conv_dim, affine=True, track_running_stats=True))
layers.append(nn.ReLU(inplace=True))
# Down-sampling layers.
curr_dim = conv_dim #这时候的64个维度
for i in range(2):
layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True, track_running_stats=True))
layers.append(nn.ReLU(inplace=True))
curr_dim = curr_dim * 2
#经过两次循环，这时 curr_dim 的维度为256
# Bottleneck layers.
for i in range(repeat_num):
layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))
# Up-sampling layers.
for i in range(2):
layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True, track_running_stats=True))
layers.append(nn.ReLU(inplace=True))
curr_dim = curr_dim // 2
#最后的维度为3维
layers.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
layers.append(nn.Tanh())
self.main = nn.Sequential(*layers)
def forward(self, x, c): #定义计算的过程
# Replicate spatially and concatenate domain information.
c = c.view(c.size(0), c.size(1), 1, 1) #view 相当于Numpy中的reshape
c = c.repeat(1, 1, x.size(2), x.size(3)) #沿着指定的维度重复tensor
x = torch.cat([x, c], dim=1) #将输入图像x,label向量c,串联
return self.main(x)
class Discriminator(nn.Module):
"""Discriminator network with PatchGAN."""
def __init__(self, image_size=128, conv_dim=64, c_dim=5, repeat_num=6):
super(Discriminator, self).__init__()
layers = []
layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
layers.append(nn.LeakyReLU(0.01))
curr_dim = conv_dim
for i in range(1, repeat_num):
layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1))
layers.append(nn.LeakyReLU(0.01))
curr_dim = curr_dim * 2
kernel_size = int(image_size / np.power(2, repeat_num))
self.main = nn.Sequential(*layers) #将层加入到神经网络
self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=3, stride=1, padding=1, bias=False)#D判读图像的真假
self.conv2 = nn.Conv2d(curr_dim, c_dim, kernel_size=kernel_size, bias=False)#判别输入图像的label.
def forward(self, x):
h = self.main(x) #这里的X表示训练时的图像，经过main()后生成2048维数据
out_src = self.conv1(h) #out_src 表示图像的真假
out_cls = self.conv2(h) # out_cls 表示图像的标签
return out_src, out_cls.view(out_cls.size(0), out_cls.size(1))
from torchvision.utils import save_image
import time
import datetime
class Multi_Solver(object):
"""Solver for training and testing StarGAN."""
def __init__(self, mul_loader):
"""Initialize configurations."""
# Data loader.
self.celeba_loader = mul_loader
# Model configurations.
self.c_dim = 20
self.c2_dim = c2_dim
self.image_size = image_size
self.g_conv_dim = g_conv_dim
self.d_conv_dim = d_conv_dim
self.g_repeat_num = g_repeat_num
self.d_repeat_num = d_repeat_num
self.lambda_cls = lambda_cls
self.lambda_rec = lambda_rec
self.lambda_gp = lambda_gp
# Training configurations.
self.dataset = dataset
self.batch_size = batch_size
self.num_iters = num_iters
self.num_iters_decay = num_iters_decay
self.g_lr = g_lr
self.d_lr = d_lr
self.n_critic = n_critic
self.beta1 = beta1
self.beta2 = beta2
self.resume_iters = resume_iters
self.selected_attrs = selected_attrs
# Test configurations.
self.test_iters = test_iters
# Miscellaneous.
self.use_tensorboard = use_tensorboard
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#self.device = torch.device('cpu')
# Directories.
self.log_dir = log_dir
self.sample_dir = sample_dir
self.model_save_dir = model_save_dir
self.result_dir = result_dir
# Step size.
self.log_step = log_step
self.sample_step = sample_step
self.model_save_step = model_save_step
self.lr_update_step = lr_update_step
# Build the model and tensorboard.
self.build_model()
if self.use_tensorboard:
self.build_tensorboard()
def build_model(self):
"""Create a generator and a discriminator."""
if self.dataset in ['CelebA', 'RaFD']:
self.G = Generator(self.g_conv_dim, self.c_dim, self.g_repeat_num)
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim, self.d_repeat_num)
elif self.dataset in ['Both']:
self.G = Generator(self.g_conv_dim, self.c_dim+self.c2_dim+2, self.g_repeat_num) # 2 for mask vector.
self.D = Discriminator(self.image_size, self.d_conv_dim, self.c_dim+self.c2_dim, self.d_repeat_num)
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 denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def build_tensorboard(self):
self.logger = Logger(self.log_dir)
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
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 update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
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 reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def classification_loss(self, logit, target, dataset='CelebA'):
"""Compute binary or softmax cross entropy loss."""
if dataset == 'CelebA':
return F.binary_cross_entropy_with_logits(logit, target, size_average=False) / logit.size(0)
elif dataset == 'RaFD':
return F.cross_entropy(logit, target)
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
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):
"""Train StarGAN within a single dataset."""
# Set data loader.
data_loader = self.celeba_loader
data_iter = iter(data_loader)
# Learning rate cache for decaying.
g_lr = self.g_lr
d_lr = self.d_lr
x_fixed, c_org,c_trg = next(data_iter)
x_fixed = x_fixed.to(self.device)
# Start training from scratch or resume training.
start_iters = 0
if self.resume_iters: #参数resume_iters 设置为none
start_iters = self.resume_iters #可以不连续训练，从之前训练好后的结果处开始
self.restore_model(self.resume_iters)
# Start training.
print('Start training...')
start_time = time.time()
for i in range(start_iters, self.num_iters):
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# Fetch real images and labels.
try:
x_real, label_org, label_trg = next(data_iter)
except:
data_iter = iter(data_loader)
x_real, label_org, label_trg= next(data_iter)
# Generate target domain labels randomly.
#rand_idx = torch.randperm(label_org.size(0)) #取出真实标签，然后在真实标签里面随机选取一个
#label_trg = label_org[rand_idx] # 真实label中任意选取一个标签
c_org = label_org.clone()
c_trg = label_trg.clone()
x_real = x_real.to(self.device) # Input images.
c_org = c_org.to(self.device) # Original domain labels.
c_trg = c_trg.to(self.device) # Target domain labels.
label_org = label_org.to(self.device) # Labels for computing classification loss.
label_trg = label_trg.to(self.device) # Labels for computing classification loss.
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
out_src, out_cls = self.D(x_real) #D接受的就只是一幅图像
d_loss_real = - torch.mean(out_src) # d_loss_real最小，那么 out_src 最大==1 （针对图像）
d_loss_cls = self.classification_loss(out_cls, label_org, self.dataset) #针对标签
x_fake = self.G(x_real, c_trg) #x_fake 生成一个图像数据
out_src, out_cls = self.D(x_fake.detach())#在这里表示不用求上面一行中的G的梯度
d_loss_fake = torch.mean(out_src) #假图像为0
# Compute loss for gradient penalty.
#计算梯度惩罚因子alpha,根据alpha结合x_real,x_fake,输入判别网络,计算梯度,得到梯度损失函数,
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
# alpha是一个随机数 tensor([[[[ 0.7610]]]])
x_hat = (alpha * x_real.data + (1 - alpha) * x_fake.data).requires_grad_(True)
# x_hat是一个图像大小的张量数据，随着alpha的改变而变化
out_src, _ = self.D(x_hat) #x_hat 表示梯度惩罚因子
d_loss_gp = self.gradient_penalty(out_src, x_hat)
d_loss = d_loss_real + d_loss_fake + self.lambda_cls * d_loss_cls + self.lambda_gp * d_loss_gp
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_cls'] = d_loss_cls.item()
loss['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
#生成网络的作用是,输入original域的图可以生成目标域的图像,输入为目标域的图像,生成original域的图像（重建）
if (i+1) % self.n_critic == 0:
# Original-to-target domain.
#将真实图像输入x_real和假的标签c_trg输入生成网络,得到生成图像x_fake
x_fake = self.G(x_real, c_trg)
out_src, out_cls = self.D(x_fake)
g_loss_fake = - torch.mean(out_src) #这里是对抗损失，希望生成的假图像为1
g_loss_cls = self.classification_loss(out_cls, label_trg, self.dataset)#向目标标签进行转化
# Target-to-original domain.
x_reconst = self.G(x_fake, c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# Backward and optimize.
g_loss = g_loss_fake + self.lambda_rec * g_loss_rec + self.lambda_cls * g_loss_cls
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_cls'] = g_loss_cls.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i+1) % self.log_step == 0:
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)
print(log)
if self.use_tensorboard:
for tag, value in loss.items():
self.logger.scalar_summary(tag, value, i+1)
# Translate fixed images for debugging. 每100轮保存一次图像
if (i+1) % self.sample_step == 0:
with torch.no_grad():
x_fake_list = [x_fixed]
x_fake_list.append(self.G(x_fixed, c_trg))
x_concat = torch.cat(x_fake_list, dim=3)
sample_path = os.path.join(self.sample_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), sample_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(sample_path))
# Save model checkpoints. 每100轮保存一下模型
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))
# Decay learning rates. 降低学习率
if (i+1) % self.lr_update_step == 0 and (i+1) > (self.num_iters - self.num_iters_decay):
g_lr -= (self.g_lr / float(self.num_iters_decay))
d_lr -= (self.d_lr / float(self.num_iters_decay))
self.update_lr(g_lr, d_lr)
print ('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
solver = Multi_Solver(mul_loader)
solver.train()
