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 = transforms.Compose([
transforms.Resize(opt.image_size), #重新设置图片大小，opt.image_size默认值为96
transforms.CenterCrop(opt.image_size), #从中心截取大小为opt.image_size的图片
transforms.ToTensor(), #转为Tensor格式，并将值取在[0,1]中
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) #标准化，得到在[-1,1]的值
])
dataset = datasets.ImageFolder(opt.data_path, transform=transforms) #从data中读取图片，图片类别会设置为文件夹名faces
dataloader = torch.utils.data.DataLoader(dataset, #然后对得到的图片进行批处理，默认一批为256张图，使用4个进程读取数据
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
drop_last=True # 什么鬼
)
# 网络，gnet为生成器，dnet为判别器
gnet, dnet = GNet(opt), DNet(opt)
map_location = lambda storage, loc: storage
if opt.dnet_path:
dnet.load_state_dict(torch.load(opt.dnet_path, map_location=map_location))
if opt.gnet_path:
gnet.load_state_dict(torch.load(opt.gnet_path, map_location=map_location))
def generate(**kwargs):#进行验证
"""
随机生成动漫头像，并根据dnet的分数选择较好的
"""
for k_, v_ in kwargs.items():
setattr(opt, k_, v_)
#device=torch.device('cuda') if opt.gpu else torch.device('cpu')
gnet, dnet = GNet(opt).eval(), DNet(opt).eval()
noises = torch.randn(opt.get_search_num, opt.nd, 1, 1).normal_(opt.noise_mean, opt.noise_std)
#noises = noises.to(device)
noises = noises.cuda()
map_location = lambda storage, loc: storage
dnet.load_state_dict(torch.load(opt.dnet_path, map_location=map_location))
gnet.load_state_dict(torch.load(opt.gnet_path, map_location=map_location))
dnet.cuda()
gnet.cuda()
# 生成图片，并计算图片在判别器的分数
fake_img = gnet(noises)
scores = dnet(fake_img).detach()
# 挑选最好的某几张，默认opt.get_num=64张，并得到其索引
indexs = scores.topk(opt.get_num)[1] # tokp()返回元组，一个为分数，一个为索引
result = []
for i in indexs:
result.append(fake_img.data[i])
# 保存图片
tv.utils.save_image(torch.stack(result), opt.get_img, normalize=True, range=(-1, 1))
def generate():
opt = Config()
criterion = nn.BCELoss()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# 训练可能多卡，预测一张就够了，所以有点小不同
#dnet = torch.load('dnet1.pth').to(device)#可能需要从其他GPU移动到0号，若满足条件则不作为
#gnet = torch.load('gnet1.pth').to(device)
dnet = DNet(opt).to(device)
gnet = GNet(opt).to(device)
state_dict = torch.load('dd.pth')
new_state_dict = OrderedDict()
for k,v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
dnet.load_state_dict(new_state_dict)
state_dict = torch.load('gg.pth')
new_state_dict = OrderedDict()
for k,v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
gnet.load_state_dict(new_state_dict)
dnet.eval()
gnet.eval()
noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
#with torch.no_grad():
fake = gnet(noise)
output = dnet(fake)
label = torch.full((opt.batch_size, ), opt.real_label, device=device)
d_err_fake = criterion(output, label) # 生成图像的损失；还是tensor
mean_score = output.mean() #生成图像的平均得分；还是tensor
fake_img = vutils.make_grid(fake, normalize=True)
writer = SummaryWriter(log_dir='generate_rusult')
writer.add_image('fake_img', fake_img)
writer.close()
print('生成图像的平均损失值：%.4f'%d_err_fake.item())
print('生成图像的平均得分：%.4f'%mean_score.item())
def generate(opt, device):
criterion = nn.BCELoss()
dnet = DNet(opt).to(device) # 可能需要从其他GPU移动到0号，若满足条件则不作为
gnet = GNet(opt).to(device)
state_dict = torch.load('dnet.pth')
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module.'
new_state_dict[name] = v
dnet.load_state_dict(new_state_dict)
state_dict = torch.load('gnet.pth')
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:]
new_state_dict[name] = v
gnet.load_state_dict(new_state_dict)
dnet.eval()
gnet.eval()
noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
with torch.no_grad():
fake = gnet(noise)
output = dnet(fake)
label = torch.full((opt.batch_size, ), opt.real_label, device=device)
d_err_fake = criterion(output, label) # 生成图像的损失；还是tensor
mean_score = output.mean() #生成图像的平均得分；还是tensor
fake_img = vutils.make_grid(fake, normalize=True)
writer = SummaryWriter(log_dir='generate_result')
writer.add_image('fake_img', fake_img)
writer.close()
print('生成图像的平均损失值：%.4f' % d_err_fake.item())
print('生成图像的平均得分：%.4f' % mean_score.item())
def train():
opt = Config() # 配置的实例
dataloader = data_loader(opt)
criterion = nn.BCELoss()
device = torch.device("cuda: 0, 1, 2" if torch.cuda.is_available() else "cpu") # 有cuda这个架构就使用0卡，无则cpu;device(可以指定任何设备，cpu,哪块或哪几块显卡)
gnet = GNet(opt).to(device)
dnet = DNet(opt).to(device)
#writer.add_graph(gnet)'''做实验试验下第二个参数'''
#writer.add_graph(dnet)
if device.type == 'cuda':
gnet = nn.DataParallel(gnet, [0, 1, 2])
dnet = nn.DataParallel(dnet, [0, 1, 2])
gnet.apply(weight_init)
dnet.apply(weight_init)
print('Generative NetWork:')
print(gnet)
print('')
print('Discriminative NetWork:')
print(dnet)
g_optimizer = optim.Adam(gnet.parameters(), lr=opt.lr1, betas=(opt.beta1, 0.999))
d_optimizer = optim.Adam(dnet.parameters(), lr=opt.lr2, betas=(opt.beta1, 0.999))
print('g_optimizer:')
print(g_optimizer)
print('d_optimizer:')
print(d_optimizer)
writer = SummaryWriter(log_dir='result_shiyan')
#dummy1_input = torch.rand(opt.batch_size, 3, 96,96)
#dummy2_input = torch.rand(opt.batch_size, opt.nd,1,1)
#writer.add_graph(dnet, dummy1_input.detach())
#writer.add_graph(gnet, dumm2_input.detach())
# Training Loop
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0
fixed_noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
print("Starting Training Loop...")
# For each epoch
for epoch in range(1, opt.max_epoch + 1):
# For each batch in the dataloader
print(len(dataloader))
print(type(dataloader))
for i, (imgs, _) in enumerate(dataloader, 1):
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
############################
## Train with all-real batch
dnet.zero_grad()
'''先训练判别器，再训练生成器'''
# Format batch
real_img = imgs.to(device)
label = torch.full((opt.batch_size, ), opt.real_label, device=device)
# Forward pass real batch through D
output = dnet(real_img) # 在model模块中已经被展成一维的啦
# Calculate loss on all-real batch
d_err_real = criterion(output, label)
# Calculate gradients for D in backward pass
d_err_real.backward()
D_x = output.mean().item() # 真实图片的平均得分，当然是接近1越好
## Train with all-fake batch
# Generate batch of latent vectors latent:隐藏的，潜伏的
noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
# Generate fake image batch with G
fake = gnet(noise)
label.fill_(opt.fake_label)
# Classify all fake batch with D
output = dnet(fake.detach())
# Calculate D's loss on the all-fake batch
d_err_fake = criterion(output, label)
# Calculate the gradients for this batch
d_err_fake.backward()
D_G_z1 = output.mean().item() # 假图像的分数，自然是越接近0越好
# Add the gradients from the all-real and all-fake batches
d_err = d_err_real + d_err_fake #tensor(1.272)+tensor(0.183)可以直接相加，不需要先取出数值。tensor
# 自成体系，tensor和tensor的加减乘除和标量一模一样；只是tensor和标量之间不能直接算
# Update D
d_optimizer.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
gnet.zero_grad()
label.fill_(opt.real_label) # fake labels are real for generator cost
# Since we just updated D, perform another forward pass of all-fake batch through D
output = dnet(fake) # 更新了一步D网络后，同样一批假图片，自然是希望判别得分output比更新前的假图片得分要小,也就是使下面的g_err扩大
# Calculate G's loss based on this output
g_err = criterion(output, label)
'''生成器就是要把假图片往真标签身上凑；所以假图片+真标签，进行比较后，损失越小越好'''
# Calculate gradients for G
g_err.backward()
D_G_z2 = output.mean().item() # 因为更新过一次判别器，所以这个假图片的output均值应该比上面的假图片的output均值更接近0才健康
# Update G
g_optimizer.step()
# Output training stats
if i % 50 == 0:
print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\treal_img_mean_score: %.4f\tfake_img_mean_score_1/2: %.4f / %.4f'
% (epoch, opt.max_epoch, i, len(dataloader),d_err.item(), g_err.item(), D_x, D_G_z1, D_G_z2))
# Save Losses for plotting later
G_losses.append(g_err.item())
D_losses.append(d_err.item())
writer.add_scalars('dnet_gnet_loss', {'G_losses': G_losses[iters], 'D_losses': D_losses[iters]}, iters)
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % 500 == 0) or ((epoch == opt.max_epoch) and (i == len(dataloader))):
with torch.no_grad():
fake = gnet(fixed_noise)#.detach().cpu()
img_list.append(vutils.make_grid(fake, normalize=True))
'''还不知道合成的图有多少个小图呢'''
writer.add_image('fake%d'%(iters/500), img_list[int(iters/500)], int(iters/500))
iters += 1
#torch.save(dnet, 'dnet1.pth')
#torch.save(gnet, 'gnet1.pth')
torch.save(dnet.state_dict(), 'dd.pth')
torch.save(gnet.state_dict(), 'gg.pth')
#writer = SummaryWriter(log_dir='result_')
print('最后的iters为: %d'%iters)
print('G_losses长度为: %d'%len(G_losses))
print('D_losses长度为: %d'%len(D_losses))
#for i in range(iters):
#writer.add_scalars('dnet_gnet_loss', {'G_losses': G_losses[i], 'D_losses': D_losses[i]}, i)
print('img_list的长度为: %d'%len(img_list))
#for i in range(len(img_list)):
#writer.add_image('fake%d'%i, img_list[i], i)
#writer.add_graph(dnet, input_to_model=(torch.rand(opt.batch_size, 3, 96, 96), ), verbose=False)
#writer.add_graph(gnet, input_to_model=(torch.rand(opt.batch_size, op.nd, 1, 1), ), verbose=False)
writer.close()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--action",
type=str,
default='test',
help="train or test")
args = parser.parse_args()
config = load_config()
# 使用tensorboard
time_now = datetime.now().isoformat()
if not os.path.exists(config.RUN_PATH):
os.mkdir(config.RUN_PATH)
writer = SummaryWriter(log_dir=os.path.join(config.RUN_PATH, time_now))
# 随机数种子
torch.manual_seed(config.SEED)
torch.cuda.manual_seed(config.SEED)
np.random.seed(config.SEED)
random.seed(config.SEED)
# INIT GPU
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU)
if torch.cuda.is_available():
config.DEVICE = torch.device("cuda")
print('\nGPU IS AVAILABLE')
torch.backends.cudnn.benchmark = True
else:
config.DEVICE = torch.device("cpu")
##############################################################
# Initialise the generator and discriminator with the UNet and
# DNet architectures respectively.
generator = UNet(True).to(config.DEVICE)
discriminator = DNet().to(config.DEVICE)
###################################################################
# Create ADAM optimizer for the generator as well the discriminator.
# Create loss criterion for calculating the L1 and adversarial loss.
d_optimizer = optim.Adam(discriminator.parameters(),
betas=(0.5, 0.999),
lr=config.LR)
g_optimizer = optim.Adam(generator.parameters(),
betas=(0.5, 0.999),
lr=config.LR)
d_loss = nn.BCELoss()
g_loss_bce = nn.BCELoss()
g_loss_l1 = nn.L1Loss()
# 加载数据集
if args.action == 'train':
train_dataset = LABDataset(config, config.TRAIN_PATH)
len_train = len(train_dataset)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=config.BATCH_SIZE, shuffle=True)
iter_per_epoch = len(train_loader)
train_(config, train_loader, generator, discriminator, d_optimizer,
g_optimizer, d_loss, g_loss_bce, g_loss_l1, len_train,
iter_per_epoch, writer)
if args.action == "test":
test_dataset = LABDataset(config, config.TEST_PATH)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False)
test(config, test_loader, generator, g_loss_l1)
import torch import torch.optim as optim import torch.nn as nn from torch.autograd import Variable from torch import cat from model import UNet, DNet import data_loader from data_loader import * ############################################################## # Initialise the generator and discriminator with the UNet and # DNet architectures respectively. generator = UNet(True) discriminator = DNet() ################################################################## # Utilize GPU for performing all the calculations performed in the # forward and backward passes. Thus allocate all the generator and # discriminator variables on the default GPU device. generator.cuda() discriminator.cuda() ################################################################### # Create ADAM optimizer for the generator as well the discriminator. # Create loss criterion for calculating the L1 and adversarial loss. d_optimizer = optim.Adam(discriminator.parameters(), betas=(0.5, 0.999), lr=0.0002) g_optimizer = optim.Adam(generator.parameters(), betas=(0.5, 0.999), lr=0.0002) d_criterion = nn.BCELoss()
parser.add_argument('--frequency',
type=int,
default=5,
help='frequency of saving model\'s parameters')
opt = parser.parse_args()
if torch.cuda.is_available() and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
Gnet_AB = GNet(opt.G_init_filter, opt.G_depth, opt.G_width)
Gnet_BA = GNet(opt.G_init_filter, opt.G_depth, opt.G_width)
Dnet_A = DNet(opt.D_init_filter, opt.D_depth)
Dnet_B = DNet(opt.D_init_filter, opt.D_depth)
if opt.cuda:
Gnet_AB.cuda()
Gnet_BA.cuda()
Dnet_A.cuda()
Dnet_B.cuda()
# Weight Initialization from a Gaussian distribution N(0, 0:02)
Gnet_AB.apply(weights_init_normal)
Gnet_BA.apply(weights_init_normal)
Dnet_A.apply(weights_init_normal)
Dnet_B.apply(weights_init_normal)
# Lossess
def train(opt, device):
dataloader = data_loader(opt)
gnet = GNet(opt).to(device)
dnet = DNet(opt).to(device)
#writer.add_graph(gnet)'''做实验试验下第二个参数'''
#writer.add_graph(dnet)
if device.type == 'cuda': # 就算device里有多个GPU可见，但是若不用分发功能，仍然只有第0块在跑
gnet = nn.DataParallel(gnet, [0, 1, 2]) # list(range(ngpu))不好使，只能用前几个
dnet = nn.DataParallel(dnet, [0, 1, 2])
gnet.apply(
weight_init) # 也就是初始化了下面的d/gnet.parameters()；不进行初始化则会系统给你进行一次随机初始
dnet.apply(weight_init)
print('Generative NetWork:')
print(gnet)
print('')
print('Discriminative NetWork:')
print(dnet)
criterion = nn.BCELoss()
'''
params (iterable): iterable of parameters to optimize or dicts defining parameter groups
除了下面的整体赋值，还可以通过迭代给优化器赋值，把模型中所有需要参数的过程都分别设置值;如学长代码：
optimizer = optim.SGD([
{'params': model.features.parameters(), 'lr': 0.1 * lr},
{'params': model.sample_128.parameters(), 'lr': lr},
{'params': model.sample_256.parameters(), 'lr': lr},
{'params': model.fc_concat.parameters(), 'lr': lr}
], lr=1e-1, momentum=0.9, weight_decay=1e-5)
'''
g_optimizer = optim.Adam(gnet.parameters(),
lr=opt.lr1,
betas=(opt.beta1, 0.999))
d_optimizer = optim.Adam(dnet.parameters(),
lr=opt.lr2,
betas=(opt.beta1, 0.999))
# 优化器只会进行一次初始赋值，其他都是反向调整
print('g_optimizer:')
print(g_optimizer)
print('d_optimizer:')
print(d_optimizer)
writer = SummaryWriter(log_dir='train_result')
# 定义writer时候就会生成events文件，而tensorboard执行时会搜索大文件下的所有路径，找出所有需要的文件
#dummy1_input = torch.rand(opt.batch_size, 3, 96,96)
#dummy2_input = torch.rand(opt.batch_size, opt.nd,1,1)
#writer.add_graph(dnet, dummy1_input)
#writer.add_graph(gnet, dumm2_input
# Training Loop
# Lists to keep track of progress
'''完全可以不用列表，但是为了以后可能有其他用，就保留了'''
img_list = []
G_losses = []
D_losses = []
iters = 0
fixed_noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
print("Starting Training Loop...")
dnet.train()
gnet.train()
# 不写也默认为train模式；当有BN层和dropout层时，肯定得考虑模式切换，因为训练时这两个层有变化，验证时不能让它变，而eval()模式就不变，train()会变
# For each epoch
for epoch in range(1, opt.max_epoch + 1):
# For each batch in the dataloader
print(len(dataloader))
print(type(dataloader))
for i, (imgs, _) in enumerate(dataloader, 1):
# torch.utils.data.DataLoader()返回的就是二元组组成的一个特殊的对象(不是列表等，也不能切片);
# 在MNIST数据集中img, label = data；这些动漫头像没有标签，打印出来后发现是tensor([0, 0,...0, 0 0])
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
############################
## Train with all-real batch
dnet.zero_grad()
'''先训练判别器，再训练生成器'''
# Format batch
real_img = imgs.to(device) # 每个batch.to(device)
# torch.full((2,3), 1.2),第一个参数必须是元组，可以是任意维数，但想要一维填充时也得为元组，而元组只有一个元素时后面必须有个,
label = torch.full((opt.batch_size, ),
opt.real_label,
device=device)
# Forward pass real batch through D
output = dnet(real_img) # 在model模块中已经被展成一维的啦
# Calculate loss on all-real batch
d_err_real = criterion(output, label) # 平均损失
# Calculate gradients for D in backward pass
d_err_real.backward()
D_x = output.mean().item() # 真实图片的平均得分，当然是接近1越好
## Train with all-fake batch
# Generate batch of latent vectors latent:隐藏的，潜伏的
noise = torch.randn(opt.batch_size, opt.nd, 1, 1, device=device)
# gnet会生成opt.batch_size个图像，因为一个(opt.nd,1,1)可以生成一个图像；在gnet中，每张图有otp.nd个feature maps
# ,每个feature map大小为1 x 1,所以每个opt.nd(也就是一个值)，控制着生成图像中的一个特征
# Generate fake image batch with G
fake = gnet(noise)
label.fill_(opt.fake_label)
# Classify all fake batch with D
output = dnet(fake.detach())
# Calculate D's loss on the all-fake batch
d_err_fake = criterion(output, label)
# Calculate the gradients for this batch
d_err_fake.backward()
D_G_z1 = output.mean().item() # 假图像的分数，自然是越接近0越好
# Add the gradients from the all-real and all-fake batches
d_err = d_err_real + d_err_fake #tensor(1.272)+tensor(0.183)可以直接相加，不需要先取出数值。
# tensor自成体系，tensor和tensor的加减乘除和标量一模一样；只是tensor和标量之间不能直接算
# Update D
d_optimizer.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
gnet.zero_grad()
label.fill_(
opt.real_label) # fake labels are real for generator cost
# Since we just updated D, perform another forward pass of all-fake batch through D
output = dnet(
fake
) # 更新了一步D网络后，同样一批假图片，自然是希望判别得分output比更新前的假图片得分要小,也就是使下面的g_err扩大
# Calculate G's loss based on this output
g_err = criterion(output, label)
'''生成器就是要把假图片往真标签身上凑；所以假图片+真标签，进行比较后，损失越小越好'''
# Calculate gradients for G
g_err.backward()
D_G_z2 = output.mean().item(
) # 因为更新过一次判别器，所以这个假图片的output均值应该比上面的假图片的output均值更接近0才健康
# Update G
g_optimizer.step()
# Output training stats
if i % 50 == 0:
print(
'[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\treal_img_mean_score: %.4f\tfake_img_mean_score_1/2: %.4f / %.4f'
% (epoch, opt.max_epoch, i, len(dataloader), d_err.item(),
g_err.item(), D_x, D_G_z1, D_G_z2))
# Save Losses for plotting later
G_losses.append(g_err.item())
D_losses.append(d_err.item())
writer.add_scalars('dnet_gnet_loss', {
'G_losses': G_losses[iters],
'D_losses': D_losses[iters]
}, iters)
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % 500 == 0) or ((epoch == opt.max_epoch) and
(i == len(dataloader))):
with torch.no_grad(
): # 上下文管理，处于with范围内的tensor待会不反向，所以前向时不用求局部梯度了，节省计算。因为forward时就会把每层对应局部梯度公式求出来
fake = gnet(
fixed_noise
) #.detach().cpu() 截断再放CPU里没什么特殊用啊，有没有效果一样，只是拷贝一份假图片存放cpu里
img_list.append(vutils.make_grid(fake, normalize=True))
'''还不知道合成的图有多少个小图呢'''
writer.add_image('fake%d' % (iters / 500),
img_list[int(iters / 500)], int(iters / 500))
iters += 1
torch.save(dnet.state_dict(), 'dnet.pth')
torch.save(gnet.state_dict(), 'gnet.pth')
writer.close()
'''