def build_model(self, inputs=None): # inputs if inputs is None: self.inputs = tf.placeholder(self.dtype, self.input_shape, name='Input') else: self.inputs = tf.identity(inputs, name='Input') self.inputs.set_shape(self.input_shape) inputs = self.inputs # convert to linear inputs = layers.Gamma2Linear(inputs, self.transfer) if self.input_range == 2: inputs = inputs * 2 - 1 # forward pass self.generator = Generator('Generator', self.config) outputs = self.generator(inputs, reuse=None) # outputs if self.output_range == 2: # outputs = tf.tanh(outputs) outputs = tf.multiply(outputs + 1, 0.5) # convert to gamma self.outputs = layers.Linear2Gamma(outputs, self.transfer) self.outputs = tf.identity(self.outputs, name='Output') self.outputs_gamma = (self.outputs if self.transfer == self.loss_transfer else layers.Linear2Gamma(outputs, self.loss_transfer)) # all the saver variables self.svars = self.generator.svars # all the restore variables self.rvars = self.generator.rvars # return outputs return self.outputs
def main(): cudnn.benchmark = True global args args = parser.parse_args() print(args) if not os.path.exists(args.model_path): os.makedirs(args.model_path) if not os.path.exists(args.sample_path): os.makedirs(args.sample_path) # Networks g_path = os.path.join(args.model_path, 'generator-%d.pkl' % (args.num_epochs)) generator = Generator() generator.load_state_dict(torch.load(g_path)) generator.eval() if torch.cuda.is_available(): generator = generator.cuda() # Sample the images noise = to_variable(torch.randn(args.sample_size, args.z_dim)) fake_images = generator(noise) sample_path = os.path.join(args.sample_path, 'fake_samples-final.png') torchvision.utils.save_image(denorm(fake_images.data), sample_path) print("Saved sampled images to '%s'" % sample_path)
def build_model(self, inputs=None, target_domains=None): # inputs if inputs is None: self.inputs = tf.placeholder(self.dtype, self.input_shape, name='Input') else: self.inputs = tf.identity(inputs, name='Input') self.inputs.set_shape(self.input_shape) # target domains if target_domains is None: self.target_domains = tf.placeholder(tf.int64, self.domain_shape, name='Domain') else: self.target_domains = tf.identity(target_domains, name='Domain') self.target_domains.set_shape(self.domain_shape) # forward pass self.generator = Generator('Generator', self.config) self.outputs = self.generator(self.inputs, self.target_domains, reuse=None) # outputs self.outputs = tf.identity(self.outputs, name='Output') # all the saver variables self.svars = self.generator.svars # all the restore variables self.rvars = self.generator.rvars # return outputs return self.outputs
def __init__(self, bsize): self.bsize = bsize self.device = 'cuda:0' self.G = Generator().cuda() self.D = NLayerDiscriminator(3).cuda() init_weights(self.G) init_weights(self.D) #self.G = torch.nn.DataParallel(self.G).cuda() #self.D = torch.nn.DataParallel(self.D).cuda() #self.zmap = torch.nn.DataParallel(self.zmap).cuda() self.optimizers = [] self.optimizer_G = torch.optim.Adam(self.G.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(self.D.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.metric = None self.reg_param = 10. def lambda_rule(iteration): if iteration < 100000: lr_l = 1 if 100000 <= iteration and iteration < 500000: lr_l = 0.5 if iteration >= 500000: lr_l = 0.5 * (1 - (iteration - 500000)/ float(1000000 - 500000 + 1)) return lr_l self.schedulers = [lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) for optimizer in self.optimizers]
def set_model(args): """Configure models and loss""" # models encoder = Encoder(out_dim=args.z_dim).cuda() generator = Generator(z_dim=args.z_dim).cuda() critic = Critic().cuda() if torch.cuda.device_count() > 1: print("Use device count: {}".format(torch.cuda.device_count())) encoder = torch.nn.DataParallel(encoder) generator = torch.nn.DataParallel(generator) critic = torch.nn.DataParallel(critic) encoder.cuda() generator.cuda() critic.cuda() cudnn.benchmark = True models = Models(encoder, generator, critic) optim_encoder = torch.optim.Adam(encoder.parameters(), lr=args.lr_encoder) optim_generator = torch.optim.Adam(generator.parameters(), lr=args.lr_generator) optim_critic = torch.optim.Adam(critic.parameters(), lr=args.lr_critic) # critieron l2_reconstruct_criterion = nn.MSELoss().cuda() return models, optim_encoder, optim_generator, optim_critic, l2_reconstruct_criterion
def __init__(self, cfg): self.translator = Generator(cfg.in_ch, cfg.out_ch, cfg.ngf, cfg.n_blocks) self.transform_fn = self.get_transform_fn(cfg.img_size) self.translator.load_state_dict( torch.load(cfg.generator_weight_path, map_location="cpu"))
def main(): test_image_dataset = image_preprocessing(args.dataset) data_loader = DataLoader(test_image_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) G = Generator() checkpoint = torch.load(args.model_path) if torch.cuda.is_available(): G = nn.DataParallel(G) G.load_state_dict(checkpoint['G_model']) G.eval() if not os.path.exists(args.save_path): os.mkdir(args.save_path) if torch.cuda.is_available(): G = G.cuda() for step, data in enumerate(data_loader): real_A = to_variable(data['A']) real_B = to_variable(data['B']) fake_B = G(real_A) batch_image = torch.cat((torch.cat((real_A, fake_B), 3), real_B), 3) for i in range(args.batch_size): torchvision.utils.save_image( denorm(batch_image[i]), args.save_path + '{result_name}_{step}.jpg'.format( result_name=args.result_name, step=step * args.batch_size + i))
def __init__(self, config): self.config = config self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.generator = Generator(config).to(self.device) self.discriminator = Discriminator(config).to(self.device) self.loss_func = nn.BCELoss() self.generator_optimizer = optim.Adam( params=self.generator.parameters(), lr=config["generator_learning_rate"], betas=config["generator_betas"]) self.discriminator_optimizer = optim.Adam( params=self.discriminator.parameters(), lr=config["discriminator_learning_rate"], betas=config["discriminator_betas"]) self.true_labels = torch.ones(config["batch_size"], dtype=torch.float32, device=self.device) self.fake_labels = torch.zeros(config["batch_size"], dtype=torch.float32, device=self.device)
def get_network(LEARNING_RATE: float, device: str): G = Generator().to(device) D = Discriminator().to(device) criterion = nn.BCELoss() d_optimizer = optim.Adam(D.parameters(), lr=LEARNING_RATE) g_optimizer = optim.Adam(G.parameters(), lr=LEARNING_RATE) return G, D, criterion, d_optimizer, g_optimizer
def __init__(self): # net = Generator(ch_style=3, ch_content=1).cuda() # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cpu") net = Generator(ch_style=3, ch_content=1).to(device) net.load_state_dict( torch.load('model_weights/latest_G.pth', map_location=torch.device('cpu'))) net.eval() self.net = net
def main(): args = set_args() encoder = Encoder(out_dim=args.z_dim) generator = Generator(z_dim=args.z_dim) encoder = load_model(args, encoder, 'encoder') generator = load_model(args, generator, 'generator') loader = set_loader(args) save_one_batch_img(args, loader, generator, encoder)
def main(): os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus test_image_dataset = image_preprocessing(opt.dataset, 'val') data_loader = DataLoader(test_image_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers) G = Generator(ResidualBlock, layer_count=9) F = Generator(ResidualBlock, layer_count=9) if torch.cuda.is_available(): G = nn.DataParallel(G) F = nn.DataParallel(F) G = G.cuda() F = F.cuda() G, F, _, _, _, _, _, _, _ = load_ckp(opt.model_path, G, F) G.eval() F.eval() if not os.path.exists(opt.save_path): os.mkdir(opt.save_path) for step, data in enumerate(tqdm(data_loader)): real_A = to_variable(data['A']) real_B = to_variable(data['B']) fake_B = G(real_A) fake_A = F(real_B) batch_image = torch.cat((torch.cat( (real_A, real_B), 3), torch.cat((fake_A, fake_B), 3)), 2) for i in range(batch_image.shape[0]): torchvision.utils.save_image( denorm(batch_image[i]), opt.save_path + '{result_name}_{step}.jpg'.format( result_name=opt.result_name, step=step * opt.batch_size + i))
def __init__(self): self.generator = Generator() self.generator.to(Params.Device) self.discriminator = Discriminator() self.discriminator.to(Params.Device) self.loss_fn = nn.BCELoss() self.optimizer_g = torch.optim.Adam(self.generator.parameters(), Params.LearningRateG, betas=(Params.Beta, 0.999)) self.optimizer_d = torch.optim.Adam(self.discriminator.parameters(), Params.LearningRateD, betas=(Params.Beta, 0.999)) self.exemplar_latent_vectors = torch.randn( (64, Params.LatentVectorSize), device=Params.Device )
def evaluate_network(args): device = torch.device('cuda' if args.gpu_no >= 0 else 'cpu') check_point = torch.load(args.check_point) network = Generator(args.unet_flag).to(device).eval() network.load_state_dict(check_point['g_state_dict']) image = imload(args.image, args.imsize, args.cropsize, args.cencrop).to(device) output = network(image) imsave(output, 'output.jpg')
def _initialise_networks(self): self.generator = Generator(final_size=self.final_size) self.generator.generate_network() self.g_optimizer = Adam(self.generator.parameters(), lr=0.001, betas=(0, 0.99)) self.discriminator = Discriminator(final_size=self.final_size) self.discriminator.generate_network() self.d_optimizer = Adam(self.discriminator.parameters(), lr=0.001, betas=(0, 0.99)) self.num_channels = min(self.generator.num_channels, self.generator.max_channels) self.upsample = [Upsample(scale_factor=2**i) for i in reversed(range(self.generator.num_blocks))]
def __init__(self, opt): super(MUNIT, self).__init__() # generators and discriminators self.gen_a = Generator(opt.ngf, opt.style_dim, opt.mlp_dim) self.gen_b = Generator(opt.ngf, opt.style_dim, opt.mlp_dim) self.dis_a = Discriminator(opt.ndf) self.dis_b = Discriminator(opt.ndf) #random style code self.s_a = torch.randn(opt.display_size, opt.style_dim, 1, 1, requires_grad=True).cuda() self.s_b = torch.randn(opt.display_size, opt.style_dim, 1, 1, requires_grad=True).cuda() #optimizers dis_params = list(self.dis_a.parameters()) + list( self.dis_b.parameters()) gen_params = list(self.gen_a.parameters()) + list( self.gen_b.parameters()) self.dis_opt = torch.optim.Adam(dis_params, lr=opt.lr, beta=opt.beta1, weight_delay=opt.weight_delay) self.gen_opt = torch.optim.Adam(gen_params, lr=opt.lr, beta=opt.beta1, weight_delay=opt.weight_delay) # nerwork weight initialization self.apply(weight_init('kaiming')) self.dis_a.apply(weight_init('gaussian')) self.dis_b.apply(weight_init('gaussian'))
def main(_): parser = argparse.ArgumentParser() parser.add_argument('--option', dest='option', type=str, default='train', help='actions: train') args = parser.parse_args() if args.option not in ['train']: print('invalid option: ', args.option) print("Please input a option: train, test, or predict") else: model = Generator(tf.Session(), configure()) getattr(model, args.option)()
def main(): # Pre-settings cudnn.benchmark = True global args args = parser.parse_args() print(args) dataset = ImageFolder(args) data_loader = data.DataLoader(dataset=dataset, batch_size=args.batchSize, shuffle=True, num_workers=2) if not os.path.exists(args.model_path): os.makedirs(args.model_path) if not os.path.exists(args.sample_path): os.makedirs(args.sample_path) g_path = os.path.join(args.model_path, 'generator-%d.pkl' % (args.num_epochs)) # Load pre-trained model generator = Generator(args.batchSize) generator.load_state_dict(torch.load(g_path)) generator.eval() if torch.cuda.is_available(): generator = generator.cuda() total_step = len(data_loader) # For Print Log for i, sample in enumerate(data_loader): AtoB = args.which_direction == 'AtoB' input_A = sample['A' if AtoB else 'B'] input_B = sample['B' if AtoB else 'A'] real_A = to_variable(input_A) fake_B = generator(real_A) real_B = to_variable(input_B) # print the log info print('Validation[%d/%d]' % (i + 1, total_step)) # save the sampled images res = torch.cat((torch.cat((real_A, fake_B), dim=3), real_B), dim=3) torchvision.utils.save_image( denorm(res.data), os.path.join(args.sample_path, 'Generated-%d.png' % (i + 1)))
def generate_image(attr): attr = np.reshape(attr, (1, 5, 1, 1)) gen = Generator(6) chainer.serializers.load_npz('./results/gen', gen) z = gen.z(1) z = chainer.Variable(z) img = gen(z, attr, alpha=1.0) path = 'static/image' if not os.path.exists(path): os.makedirs(path) rand = np.random.randint(0, 100000) for im in img: filename = os.path.join(path, 'fake{}.jpg'.format(rand)) save_image(im, filename) return filename
def init_net(depth, dropout, window, cgan): input_shape = (1 if not cgan else 2, window) # Create the 3 networks gen = Generator(depth, dropout, verbose=0) discr = Discriminator(depth, dropout, input_shape, verbose=0) if not cgan else ConditionalDiscriminator( depth, dropout, input_shape, verbose=0) ae = AutoEncoder(depth, dropout, verbose=0) # Put them on cuda if available if torch.cuda.is_available(): gen.cuda() discr.cuda() ae.cuda() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("Using : " + str(device)) print("Network initialized\n") return gen, discr, ae, device
def main(): G = Generator(args.dim_disc + args.dim_cont) D = Discriminator() if os.path.isfile(args.model): model = torch.load(args.model) G.load_state_dict(model[0]) D.load_state_dict(model[1]) if use_cuda: G.cuda() D.cuda() if args.mode == "train": G, D = train(G, D) if args.model: torch.save([G.state_dict(), D.state_dict()], args.model, pickle_protocol=4) elif args.mode == "gen": gen(G)
def _init_network(self): def init_weights(m): if type(m) == nn.Conv2d: nn.init.normal_(m.weight.data, 0.0, 0.02) if type(m) == nn.BatchNorm2d: nn.init.normal_(m.weight.data, 0.0, 0.02) nn.init.constant_(m.bias.data, 0.0) generator = Generator(DCGAN_Config).to(self.device) discriminator = Discriminator(DCGAN_Config).to(self.device) if self.device.type =='cuda' and DCGAN_Config['ngpu']>1: generator = nn.DataParallel(generator, list(range(DCGAN_Config['ngpu']))) discriminator = nn.DataParallel(discriminator, list(range(DCGAN_Config['ngpu']))) generator.apply(init_weights) discriminator.apply(init_weights) print(generator) print(discriminator) return generator, discriminator
def test(): # Roading the network Network = Generator() param = torch.load(opt.test_model_name) Network.load_state_dict(param) Network = Network.to(device) Network.eval() # Roading the dataset testloader = DataLoader( TestDataset(opt.test_dataset), batch_size=1, shuffle=False, num_workers=opt.n_cpu ) with torch.no_grad(): for i, imgs in enumerate(testloader): imgs = imgs.to(device) outputs = Network(imgs) print(outputs[0].cpu().numpy().transpose(1, 2, 0).shape) # Save the images cv2.imwrite("results/" + opt.network_name + "/" + str(i) + ".png", cv2.cvtColor(outputs[0].cpu().numpy().transpose(1, 2, 0), cv2.COLOR_BGR2RGB) * 255)
def generate(savedir, model_path, num): nz = 16 width = 28 height = 28 channel = 1 device = 'cuda' model = Generator(nz, width, height, channel) model = nn.DataParallel(model) model.module.load_state_dict(torch.load(model_path)) model.eval() # 推論モードへ切り替え(Dropoutなどの挙動に影響) # 保存先のファイルを作成 if os.path.exists(savedir): n = 1 while 1: if os.path.exists('{}({})'.format(savedir, n)): n += 1 else: savedir = '{}({})'.format(savedir, n) break os.makedirs(savedir, exist_ok=True) for i in range(num): # 入力の乱数を作成 z = torch.randn(1, nz, 1, 1) z = z.to(device) gene_img = model(z) # ジェネレータの出力画像を保存 torchvision.utils.save_image(gene_img, "{}/{:04}.png".format(savedir, i + 1))
# Create directories for images, tensorboard results and saved models. if not args.dry_run: if not os.path.exists(EXPERIMENT_DIR): os.makedirs(EXPERIMENT_DIR) # Set up root experiment directory. os.makedirs(args.save_image_dir) os.makedirs(args.tensorboard_dir) os.makedirs(args.save_model_dir) WRITER = SummaryWriter(args.tensorboard_dir) # Set up TensorBoard. else: print('Dry run! Just for testing, data is not saved') DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Set up the GAN. critic_model = Critic().to(DEVICE) generator_model = Generator().to(DEVICE) # Load pre-trained models if they are provided. if args.load_critic_model_path: critic_model.load_state_dict(torch.load(args.load_critic_model_path)) if args.load_generator_model_path: generator_model.load_state_dict(torch.load(args.load_generator_model_path)) # Set up Adam optimizers for both models. critic_optimizer = optim.Adam(critic_model.parameters(), lr=args.learning_rate, betas=(0, 0.9)) generator_optimizer = optim.Adam(generator_model.parameters(), lr=args.learning_rate, betas=(0, 0.9))
transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]), mode='train') dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchsize, shuffle=True, num_workers=opt.num_workers) # print(dataloader.__len__()) # img=dataset.__getitem__(10) # print(img['A'].size()) netG_A2B = Generator(opt.input_nc, opt.output_nc, opt.ngf) netG_B2A = Generator(opt.input_nc, opt.output_nc, opt.ngf) netD_A = Discriminator(opt.input_nc, opt.ndf) netD_B = Discriminator(opt.input_nc, opt.ndf) summary(netG_A2B, input_size=(3, 256, 256), device='cpu') summary(netD_A, input_size=(3, 256, 256), device='cpu') #instialize weights netG_A2B.apply(weights_init) netG_B2A.apply(weights_init) netD_A.apply(weights_init) netD_B.apply(weights_init) #print(netG_A2B)
np.concatenate([ in_resized[i] * 255, o * 255, out_sample_images[i] * 255 ], axis=1)) for i, o in enumerate(preds) ] }, commit=False) from network import Generator, Discriminator, get_gan_network shape = (config.input_width, config.input_height, 3) image_shape = (config.output_width, config.output_height, 3) generator = Generator(shape).generator() discriminator = Discriminator(image_shape).discriminator() #adam = optimizers.Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08) generator.compile(loss="mse", optimizer="adam") discriminator.compile(loss="binary_crossentropy", optimizer="adam") gan = get_gan_network(discriminator, shape, generator, "adam", gan_metric=perceptual_distance) for e in range(1, config.num_epochs + 1): print('-' * 15, 'Epoch %d' % e, '-' * 15) for _ in range(config.steps_per_epoch):
def train(savedir, train_list_A, train_list_B, test_list_A, test_list_B, root, epochs, batch_size): # 画像のチャンネル数 channel = 1 # LossにおけるCycle Lossの割合を決めるパラメータ(Cycle Lossにかかる係数) cycle_rate = 10.0 # LossにおけるIdentity Lossの割合を決めるパラメータ(Identity Lossにかかる係数) iden_rate = 0 # ジェネレータのAdam設定(default: lr=0.001, betas=(0.9, 0.999), weight_decay=0) G_opt_para = {'lr': 0.0002, 'betas': (0.5, 0.9), 'weight_decay': 0} # ディスクリミネータのAdam設定 D_A_opt_para = {'lr': 0.0002, 'betas': (0.5, 0.9), 'weight_decay': 0} D_B_opt_para = {'lr': 0.0002, 'betas': (0.5, 0.9), 'weight_decay': 0} device = 'cuda' myloss = MyLoss() # 保存先のファイルを作成 if os.path.exists(savedir): num = 1 while 1: if os.path.exists('{}({})'.format(savedir, num)): num += 1 else: savedir = '{}({})'.format(savedir, num) break os.makedirs(savedir, exist_ok=True) os.makedirs('{}/generating_image'.format(savedir), exist_ok=True) os.makedirs('{}/model'.format(savedir), exist_ok=True) os.makedirs('{}/loss'.format(savedir), exist_ok=True) G_A2B_model, G_B2A_model, D_A_model, D_B_model = Generator(channel), Generator(channel), Discriminator(channel), Discriminator(channel) G_A2B_model, G_B2A_model, D_A_model, D_B_model = nn.DataParallel(G_A2B_model), nn.DataParallel(G_B2A_model), nn.DataParallel(D_A_model), nn.DataParallel(D_B_model) G_A2B_model, G_B2A_model, D_A_model, D_B_model = G_A2B_model.to(device), G_B2A_model.to(device), D_A_model.to(device), D_B_model.to(device) # 最適化アルゴリズムの設定 G_para = torch.optim.Adam(chain(G_A2B_model.parameters(), G_B2A_model.parameters()), lr=G_opt_para['lr'], betas=G_opt_para['betas'], weight_decay=G_opt_para['weight_decay']) D_A_para = torch.optim.Adam(D_A_model.parameters(), lr=D_A_opt_para['lr'], betas=D_A_opt_para['betas'], weight_decay=D_A_opt_para['weight_decay']) D_B_para = torch.optim.Adam(D_B_model.parameters(), lr=D_B_opt_para['lr'], betas=D_B_opt_para['betas'], weight_decay=D_B_opt_para['weight_decay']) # ロスの推移を保存するためのリストを確保 result = {} result['G_log_loss'] = [] result['D_A_log_loss'] = [] result['D_B_log_loss'] = [] df_A = pd.read_csv(train_list_A, usecols=['Path']) df_B = pd.read_csv(train_list_B, usecols=['Path']) df_test_A = pd.read_csv(test_list_A, usecols=['Path']) df_test_A = df_test_A.sample(frac=1) df_test_B = pd.read_csv(test_list_B, usecols=['Path']) df_test_B = df_test_B.sample(frac=1) train_dataset = LoadDataset(df_A, df_B, root, transform=Trans()) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True) test_dataset = LoadDataset(df_test_A[0:batch_size], df_test_B[0:batch_size], root, transform=Trans()) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=True) output_env('{}/env.txt'.format(savedir), batch_size, G_opt_para, D_A_opt_para, D_B_opt_para, G_A2B_model, G_B2A_model, D_A_model, D_B_model) for epoch in range(epochs): print('########## epoch : {}/{} ##########'.format(epoch+1, epochs)) G_log_loss, D_A_log_loss, D_B_log_loss = [], [], [] for img_A, img_B in tqdm(train_loader): # GPU用の変数に変換 img_A = img_A.to(device) img_B = img_B.to(device) # 真正画像をジェネレータに入力 fake_img_A = G_B2A_model(img_B) fake_img_B = G_A2B_model(img_A) # 生成画像をジェネレータに入力 rec_img_A = G_B2A_model(fake_img_B) rec_img_B = G_A2B_model(fake_img_A) # ディスクリミネータに真正画像と生成画像を入力 real_pred_A = D_A_model(img_A) real_pred_B = D_B_model(img_B) fake_pred_A = D_A_model(fake_img_A) fake_pred_B = D_B_model(fake_img_B) # ジェネレータに出力画像と同一ドメインの画像を入力(恒等写像) if iden_rate == 0: iden_img_A = None iden_img_B = None else: iden_img_A = G_B2A_model(img_A) iden_img_B = G_A2B_model(img_B) # ジェネレータのロス計算 G_loss = myloss.G_loss(fake_pred_A, fake_pred_B, torch.tensor(1.0).expand_as(fake_pred_A).to(device), img_A, img_B, rec_img_A, rec_img_B, iden_img_A, iden_img_B, alpha=cycle_rate, beta=iden_rate) G_log_loss.append(G_loss.item()) # ディスクリミネータのロス計算 D_A_loss = myloss.D_A_loss(real_pred_A, torch.tensor(1.0).expand_as(real_pred_A).to(device), fake_pred_A, torch.tensor(0.0).expand_as(fake_pred_A).to(device)) D_A_log_loss.append(D_A_loss.item()) D_B_loss = myloss.D_B_loss(real_pred_B, torch.tensor(1.0).expand_as(real_pred_B).to(device), fake_pred_B, torch.tensor(0.0).expand_as(fake_pred_B).to(device)) D_B_log_loss.append(D_B_loss.item()) # ジェネレータの重み更新 G_para.zero_grad() G_loss.backward(retain_graph=True) G_para.step() # ディスクリミネータの重み更新 D_A_para.zero_grad() D_A_loss.backward(retain_graph=True) D_A_para.step() D_B_para.zero_grad() D_B_loss.backward() D_B_para.step() result['G_log_loss'].append(statistics.mean(G_log_loss)) result['D_A_log_loss'].append(statistics.mean(D_A_log_loss)) result['D_B_log_loss'].append(statistics.mean(D_B_log_loss)) print('G_loss = {} , D_A_loss = {} , D_B_loss = {}'.format(result['G_log_loss'][-1], result['D_A_log_loss'][-1], result['D_B_log_loss'][-1])) # ロスのログを保存 with open('{}/loss/log.txt'.format(savedir), mode='a') as f: f.write('##### Epoch {:03} #####\n'.format(epoch+1)) f.write('G: {}, D_A: {}, D_B: {}\n'.format(result['G_log_loss'][-1], result['D_A_log_loss'][-1], result['D_B_log_loss'][-1])) # 定めた保存周期ごとにモデル,出力画像を保存する if (epoch+1)%10 == 0: # モデルの保存 torch.save(G_A2B_model.module.state_dict(), '{}/model/G_A2B_model_{}.pth'.format(savedir, epoch+1)) torch.save(G_B2A_model.module.state_dict(), '{}/model/G_B2A_model_{}.pth'.format(savedir, epoch+1)) torch.save(D_A_model.module.state_dict(), '{}/model/D_A_model_{}.pth'.format(savedir, epoch+1)) torch.save(D_B_model.module.state_dict(), '{}/model/D_B_model_{}.pth'.format(savedir, epoch+1)) G_A2B_model.eval() G_B2A_model.eval() # メモリ節約のためパラメータの保存は止める(テスト時にパラメータの保存は不要) with torch.no_grad(): for test_img_A, test_img_B in test_loader: fake_img_test_A = G_B2A_model(test_img_B) fake_img_test_B = G_A2B_model(test_img_A) torchvision.utils.save_image(fake_img_test_A[:batch_size], "{}/generating_image/A_epoch_{:03}.png".format(savedir, epoch+1)) torchvision.utils.save_image(fake_img_test_B[:batch_size], "{}/generating_image/B_epoch_{:03}.png".format(savedir, epoch+1)) G_A2B_model.train() G_B2A_model.train() # 定めた保存周期ごとにロスを保存する if (epoch+1)%50 == 0: x = np.linspace(1, epoch+1, epoch+1, dtype='int') plot(result['G_log_loss'], result['D_A_log_loss'], result['D_B_log_loss'], x, savedir) # 最後のエポックが保存周期でない場合に,保存する if (epoch+1)%10 != 0 and epoch+1 == epochs: torch.save(G_A2B_model.module.state_dict(), '{}/model/G_A2B_model_{}.pth'.format(savedir, epoch+1)) torch.save(G_B2A_model.module.state_dict(), '{}/model/G_B2A_model_{}.pth'.format(savedir, epoch+1)) torch.save(D_A_model.module.state_dict(), '{}/model/D_A_model_{}.pth'.format(savedir, epoch+1)) torch.save(D_B_model.module.state_dict(), '{}/model/D_B_model_{}.pth'.format(savedir, epoch+1)) G_A2B_model.eval() G_B2A_model.eval() # メモリ節約のためパラメータの保存は止める(テスト時にパラメータの保存は不要) with torch.no_grad(): for test_img_A, test_img_B in test_loader: fake_img_test_A = G_B2A_model(test_img_B) fake_img_test_B = G_A2B_model(test_img_A) torchvision.utils.save_image(fake_img_test_A[:batch_size], "{}/generating_image/A_epoch_{:03}.png".format(savedir, epoch+1)) torchvision.utils.save_image(fake_img_test_B[:batch_size], "{}/generating_image/B_epoch_{:03}.png".format(savedir, epoch+1)) G_A2B_model.train() G_B2A_model.train() x = np.linspace(1, epoch+1, epoch+1, dtype='int') plot(result['G_log_loss'], result['D_A_log_loss'], result['D_B_log_loss'], x, savedir)
if __name__ == "__main__": img_list = open(config.train['img_list'], 'r').read().split('\n') img_list.pop() #input dataloader = torch.utils.data.DataLoader( TrainDataset(img_list), batch_size=config.train['batch_size'], shuffle=True, num_workers=8, pin_memory=True) G = torch.nn.DataParallel( Generator(zdim=config.G['zdim'], use_batchnorm=config.G['use_batchnorm'], use_residual_block=config.G['use_residual_block'], num_classes=config.G['num_classes'])).cuda() D = torch.nn.DataParallel( Discriminator(use_batchnorm=config.D['use_batchnorm'])).cuda() optimizer_G = torch.optim.Adam(filter(lambda p: p.requires_grad, G.parameters()), lr=1e-4) optimizer_D = torch.optim.Adam(filter(lambda p: p.requires_grad, D.parameters()), lr=1e-4) last_epoch = -1 if config.train['resume_model'] is not None: e1 = resume_model(G, config.train['resume_model']) e2 = resume_model(D, config.train['resume_model']) assert e1 == e2 last_epoch = e1
def __init__(self, depth): self.depth = depth self.gen = Generator(self.depth) serializers.load_npz('./pg_gen.npz', self.gen)