gan.zip_GAN网络_GaN_深度学习

from __future__ import print_function import torchvision.datasets as dsets import numpy as np import time from torch.autograd import Variable from torchvision import transforms,utils import torch import torch.nn as nn import os import numpy as np from glob import glob from PIL import Image from tqdm import trange from itertools import chain import torch.optim as optim import torchvision.utils as vutils from torch.autograd import Variable import torch.nn.functional as F from torch.autograd import Variable from torch.utils.data import TensorDataset, DataLoader from tensorboardX import SummaryWriter class Dataset(object): def __init__(self,x0,x1,label,root,scale_size): self.root=root self.x0=x0 self.x1=x1 self.label=torch.from_numpy(label) self.size=label.shape[0] self.paths=glob(os.path.join(root,'pants/*')) self.shape=list(Image.open(self.paths[0]).size)+[3] self.transforms=transforms.Compose([ transforms.Scale(scale_size), transforms.ToTensor(), transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)) ]) def __getitem__(self,index): return (self.transforms(self.x0[index]), self.transforms(self.x1[index]), self.label[index]) def __len__(self): return self.size def create_pairs(root,label_path): x0_data=[] x1_data=[] label=[] a=np.loadtxt(label_path) n=a.size for i in range(n): if a[i]==1: pants_path=os.path.join(root,"pants/{}.jpg".format(i)) tops_path=os.path.join(root,"tops/{}.jpg".format(i)) x0_data.append(Image.open(pants_path).convert('RGB')) x1_data.append(Image.open(tops_path).convert('RGB')) label.append(0) if a[i]==0: pants_path=os.path.join(root,"pants/{}.jpg".format(i)) tops_path=os.path.join(root,"tops/{}.jpg".format(i)) x0_data.append(Image.open(pants_path).convert('RGB')) x1_data.append(Image.open(tops_path).convert('RGB')) label.append(1) label=np.array(label,dtype=np.int32) return x0_data,x1_data,label def create_iterator(root,label_path,scale_size,shuffle=False): x0,x1,label=create_pairs(root,label_path) ret=Dataset(x0,x1,label,root,scale_size) return ret def get_loader(root,label_path,batch_size,scale_size,num_workers=2,shuffle=True): data_iter=create_iterator(root,label_path,scale_size) train_loader=torch.utils.data.DataLoader(dataset=data_iter,batch_size=batch_size,shuffle=True, num_workers=num_workers) train_loader.shape=data_iter.shape return train_loader class GeneratorCNN(nn.Module): def __init__(self, input_channel, output_channel, conv_dims, deconv_dims, num_gpu): super(GeneratorCNN, self).__init__() self.num_gpu = num_gpu self.layers = [] prev_dim = conv_dims[0] self.layers.append(nn.Conv2d(input_channel, prev_dim, 4, 2, 1, bias=False)) self.layers.append(nn.LeakyReLU(0.2, inplace=True)) for out_dim in conv_dims[1:]: self.layers.append(nn.Conv2d(prev_dim, out_dim, 4, 2, 1, bias=False)) self.layers.append(nn.BatchNorm2d(out_dim)) self.layers.append(nn.LeakyReLU(0.2, inplace=True)) prev_dim = out_dim for out_dim in deconv_dims: self.layers.append(nn.ConvTranspose2d(prev_dim, out_dim, 4, 2, 1, bias=False)) self.layers.append(nn.BatchNorm2d(out_dim)) self.layers.append(nn.ReLU(True)) prev_dim = out_dim self.layers.append(nn.ConvTranspose2d(prev_dim, output_channel, 4, 2, 1, bias=False)) self.layers.append(nn.Tanh()) self.layer_module = nn.ModuleList(self.layers) def main(self, x): out = x for layer in self.layer_module: out = layer(out) return out def forward(self, x): return self.main(x) class DiscriminatorCNN(nn.Module): def __init__(self, input_channel, output_channel, hidden_dims, num_gpu): super(DiscriminatorCNN, self).__init__() self.num_gpu = num_gpu self.layers = [] prev_dim = hidden_dims[0] self.layers.append(nn.Conv2d(input_channel, prev_dim, 4, 2, 1, bias=False)) self.layers.append(nn.LeakyReLU(0.2, inplace=True)) for out_dim in hidden_dims[1:]: self.layers.append(nn.Conv2d(prev_dim, out_dim, 4, 2, 1, bias=False)) self.layers.append(nn.BatchNorm2d(out_dim)) self.layers.append(nn.LeakyReLU(0.2, inplace=True)) prev_dim = out_dim self.layers.append(nn.Conv2d(prev_dim, output_channel, 4, 1, 0, bias=False)) self.layers.append(nn.Sigmoid()) self.layer_module = nn.ModuleList(self.layers) def main(self, x): out = x for layer in self.layer_module: out = layer(out) return out.view(out.size(0), -1) def forward(self, x): return self.main(x) class SiameseNetwork(nn.Module): def __init__(self,num_gpu): super(SiameseNetwork, self).__init__() self.num_pgu=num_gpu self.cnn1=nn.Sequential( nn.Conv2d(3, 64, 4, 2, 1, bias=False), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, 64 * 2, 4, 2, 1, bias=False), nn.BatchNorm2d(64 * 2), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64 * 2, 64 * 4, 4, 2, 1, bias=False), nn.BatchNorm2d(64 * 4), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64 * 4, 64 * 8, 4, 2, 1, bias=False), nn.BatchNorm2d(64 * 8), nn.LeakyReLU(0.2, inplace=True), nn.ConvTranspose2d(64 * 8, 64 * 4, 4, 2, 1, bias=False), nn.BatchNorm2d(64 * 4), nn.ReLU(True), nn.ConvTranspose2d(64 * 4, 64 * 2, 4, 2, 1, bias=False), nn.BatchNorm2d(64 * 2), nn.ReLU(True), nn.ConvTranspose2d(64 * 2, 64, 4, 2, 1, bias=False), nn.BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d( 64, 3, 4, 2, 1, bias=False), nn.Sigmoid() ) def forward_once(self,x): output0=self.cnn1(x) output1=output0.view(output0.size()[0],-1) return output0,output1 def forward(self, input1, input2): output_1,outputA = self.forward_once(input1) output_2,outputB = self.forward_once(input2) return output_1, output_2,outputA,outputB class ContrastiveLoss(torch.nn.Module): """ Contrastive loss function. Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf """ def __init__(self, margin=2.0): super(ContrastiveLoss, self).__init__() self.margin = margin def forward(self, output1, output2, label): euclidean_distance = F.pairwise_distance(output1, output2) loss_contrastive = torch.mean((1-label) * torch.pow(euclidean_distance, 2) + (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2)) return loss_contrastive def weights_init(m): classname=m.__class__.__name__ if classname.find('Conv')!=-1: m.weight.data.normal_(0.0,0.02) elif classname.find('BatchNorm')!=-1: m.weight.data.normal_(1.0,0.02) m.bias.data.fill_(0) class Trainer(object): def __init__(self,train_loader,num_gpu): self.train_loader=train_loader self.num_gpu=num_gpu self.lr=0.0002 self.beta1=0.5 self.beta2=0.999 self.optimizer='adam' self.batch_size=64 self.weight_decay=0.0001 self.model_dir='/home/liuchong/C_file/model_dir' self.pic_dir='/home/liuchong/C_file/pic_dir' self.load_path='/home/liuchong/C_file/model_dir' self.start_step=0 self.log_step=10 self.ma