a GAN using Wasserstein loss and resnet to generate anime pics.
一個resnet-WGAN用于生成各種二次元頭像（你也可以使用別的影像資料集，用于生成圖片）
@本專案用于深度學習中的學習交流，如有任何問題，歡迎聯系我！聯系方式QQ：741533684
#我使用了殘差模塊設計了了兩個相對對稱的殘差網路，分別做生成對抗網路的的生成器與判別器，基本原理其實與DCGAN類似，在此基礎上，使用了不同于Binary cross entropy loss的Wasserstein loss， 并將優化器從Adam修改為RMSprop（注意：Adam會導致訓練不穩定，建議使用RMSprop或者SGD，且學習率不能太大，最好使用學習率衰減，）
之后我會上傳我訓練的模型，以供大家使用作為預訓練模型，
train.py檔案重要代碼如下：
def weights_init(m): # 初始化模型權重
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find('BatchNorm') != -1:
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=64)
parser.add_argument('--imageSize', type=int, default=96)
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--epoch', type=int, default=500, help='number of epochs to train for')
parser.add_argument('--lrd', type=float, default=5e-5,
help="Discriminator's learning rate, default=0.00005") # Discriminator's learning rate
parser.add_argument('--lrg', type=float, default=5e-5,
help="Generator's learning rate, default=0.00005") # Generator's learning rate
parser.add_argument('--data_path', default='data/', help='folder to train data') # 將資料集放在此處
parser.add_argument('--outf', default='imgv3/',
help='folder to output images and model checkpoints') # 輸出生成圖片以及保存模型的位置
opt = parser.parse_args()
# 定義是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 影像讀入與預處理
transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(opt.imageSize),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])
dataset = torchvision.datasets.ImageFolder(opt.data_path, transform=transforms)
dataloader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=opt.batchSize,
shuffle=True,
drop_last=True,
)
netG = NetG().to(device)
netG.apply(weights_init)
print('Generator:' )
print(sum(p.numel() for p in netG.parameters()))
netD = NetD().to(device)
netD.apply(weights_init)
print('Discriminator')
print(sum(p.numel() for p in netD.parameters()))
print(dataset)
netG.load_state_dict(torch.load('imgv2.5/netG_0280.pth', map_location=device)) # 這兩句用來讀取預訓練模型
netD.load_state_dict(torch.load('imgv2.5/netD_0280.pth', map_location=device)) # 這兩句用來讀取預訓練模型
criterionG = Hinge()
optimizerG = torch.optim.RMSprop(netG.parameters(), lr=opt.lrg)
optimizerD = torch.optim.RMSprop(netD.parameters(), lr=opt.lrd)
lrd_scheduler = torch.optim.lr_scheduler.StepLR(optimizerD, step_size=5, gamma=0.92)
lrg_scheduler = torch.optim.lr_scheduler.StepLR(optimizerG, step_size=5, gamma=0.92)
criterionD = Hinge()
criterion = torch.nn.BCELoss()
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0
total_lossD = 0.0
total_lossG = 0.0
label = label.unsqueeze(1)
start_epoch = 280 # 設定初始epoch大小
for epoch in range(start_epoch + 1, opt.epoch + 1):
with tqdm(total=len(dataloader), desc=f'Epoch {epoch}/{opt.epoch}', postfix=dict, mininterval=0.3) as pbar:
for i, (imgs, _) in enumerate(dataloader):
# 固定生成器G，訓練鑒別器D
# 讓D盡可能的把真圖片判別為1
imgs = imgs.to(device)
# for k in range(1,5):
outputreal = netD(imgs)
optimizerD.zero_grad()
## 讓D盡可能把假圖片判別為0
# label.data.fill_(fake_label)
noise = torch.randn(opt.batchSize, opt.nz)
# noise = torch.randn(opt.batchSize, opt.nz)
noise = noise.to(device)
fake = netG(noise) # 生成假圖
outputfake = netD(fake.detach()) # 避免梯度傳到G，因為G不用更新
lossD = criterionD(outputreal, outputfake)
total_lossD += lossD.item()
lossD.backward()
optimizerD.step()
# 固定鑒別器D，訓練生成器G
optimizerG.zero_grad()
# 讓D盡可能把G生成的假圖判別為1
output = netD(fake)
lossG = criterionG(output)
total_lossG += lossG.item()
lossG.backward()
optimizerG.step()
# print('[%d/%d][%d/%d] Loss_D: %.3f Loss_G %.3f'% (epoch, opt.epoch, i, len(dataloader), lossD.item(), lossG.item()))
pbar.set_postfix(**{'total_lossD': total_lossD / (i + 1),
'lrd':get_lr(optimizerD), 'total_lossG': total_lossG / (i + 1), 'lrg': get_lr(optimizerG)})
pbar.update(1)
lrg_scheduler.step()
lrd_scheduler.step()
vutils.save_image(fake.data,
'%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch),
normalize=True)
log = open("./log.txt", 'a')
print('[%d/%d] total_Loss_D: %.3f total_Loss_G %.3f' % (
epoch, opt.epoch, total_lossD / (len(dataloader)), total_lossG / (len(dataloader))),
file=log)
total_lossG = 0.0
total_lossD = 0.0
log.close()
if epoch % 5 == 0: # 每5個epoch，保存一次模型引數.
torch.save(netG.state_dict(), '%s/netG_%04d.pth' % (opt.outf, epoch))
torch.save(netD.state_dict(), '%s/netD_%04d.pth' % (opt.outf, epoch))以下是殘差模塊，residual block的定義：
class BasicBlock(nn.Module):
def __init__(self, in1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in1, in1 * 2, kernel_size=1,
stride=1, padding=0, bias=False)
self.bn1 =nn.BatchNorm2d(in1*2)
self.relu1 = nn.LeakyReLU(0.2)
self.conv2 = nn.Conv2d(in1*2, in1, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn2 =nn.BatchNorm2d(in1)
self.relu2 = nn.LeakyReLU(0.2)
def forward(self, x):
residual = x
out = self.conv1(x)
# out = self.bn1(out)
out = self.relu1(out)
out = self.conv2(out)
# out = self.bn2(out)
out = self.relu2(out)
out += residual
return out損失函式使用了wasserstein loss，相比于BCEloss（JS距離），能準確衡量生成器產生圖片的質量，而Hinge loss相對于W loss來說，能解決其梯度爆炸導致訓練不穩定的問題，
class Wasserstein(nn.Module):
def forward(self, pred_real, pred_fake=None):
if pred_fake is not None:
loss_real = -pred_real.mean()
loss_fake = pred_fake.mean()
loss = loss_real + loss_fake
return loss
else:
loss = -pred_real.mean()
return loss
class Hinge(nn.Module):#與Wasserstein相比，Hinge能防止梯度暴增，
def forward(self, pred_real, pred_fake=None):
if pred_fake is not None:
loss_real = F.relu(1 - pred_real).mean()
loss_fake = F.relu(1 + pred_fake).mean()
return loss_real + loss_fake
else:
loss = -pred_real.mean()
return loss
所使用的判別器代碼：
class RestNet18(nn.Module):
def __init__(self):
super(RestNet18, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3 ,stride=1, padding=1)
self.layer1 = nn.Sequential(
BasicBlock(64),
nn.AvgPool2d(3, 2),
BasicBlock(64),
BasicBlock(64),
)
self.layer2 = nn.Sequential(
nn.AvgPool2d(3,2),
BasicBlock(64),
BasicBlock(64),
)
self.layer3 = nn.Sequential(
nn.AvgPool2d(3, 2),
BasicBlock(64),
BasicBlock(64),
)
self.layer4 = nn.Sequential(
nn.AvgPool2d(3, 2),
BasicBlock(64),
BasicBlock(64)
# nn.LayerNorm([64,5,5]),
)
self.layer5 = nn.Sequential(
nn.BatchNorm2d(64),
# nn.LayerNorm([64,5,5]),
nn.ReLU(True)
)
self.fc = nn.Sequential(
nn.Linear(1600, 1),
)
def forward(self, x):
out = self.conv1(x)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = self.layer5(out)
out = torch.flatten(out,start_dim=1)
out = self.fc(out)
out = F.sigmoid(out)
return out
生成器代碼如下：
class Generator(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(nz, 64*3*3)
self.layer1 = nn.Sequential(
BasicBlock(64),
nn.UpsamplingNearest2d(scale_factor=2),
BasicBlock(64),
nn.UpsamplingNearest2d(scale_factor=2),
)
self.layer2 = nn.Sequential(
BasicBlock(64),
nn.UpsamplingNearest2d(scale_factor=2),
BasicBlock(64),
nn.UpsamplingNearest2d(scale_factor=2)
)
self.layer3 = nn.Sequential(
BasicBlock(64),
BasicBlock(64),
nn.UpsamplingNearest2d(scale_factor=2)
)
self.layer4 = nn.Sequential(
BasicBlock(64),
)
self.Conv = nn.Sequential(
BasicBlock(64),
nn.BatchNorm2d(64),
# nn.LayerNorm([64,96,96]),
nn.ReLU(True),
nn.Conv2d(64, 3, kernel_size=3, padding=1, stride=1),
nn.Tanh()
)
def forward(self, z):
x = self.linear(z)
x = x.view(batch_size,64,3,3)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.Conv(x)
return x
所有代碼在我的Github中已經上傳：
https://github.com/rabbitdeng/anime-WGAN-resnet-pytorch
readme.md中有所使用的資料集的百度云盤鏈接！
轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/291774.html
標籤：AI
上一篇：基于DCGAN的動漫頭像生成