train.py

#!/usr/bin/python3

import argparse
import itertools

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch.autograd import Variable
from PIL import Image
import torch

from model import Generator, GlobalGenerator2, InceptionV3
# from utils import ReplayBuffer
from utils import LambdaLR
from utils import channel2width
from utils import weights_init_normal
from utils import createNRandompatches
from dataset import UnpairedDepthDataset
import utils_pl
from collections import OrderedDict
import util.util as util
import networks
import numpy as np
import os

parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, help='name of this experiment')
parser.add_argument('--checkpoints_dir', type=str, default='checkpoints', help='Where checkpoints are saved')
parser.add_argument('--epoch', type=int, default=0, help='starting epoch')
parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training') # DT: changed from 100
parser.add_argument('--batchSize', type=int, default=4, help='size of the batches')
parser.add_argument('--cuda', action='store_true', help='use GPU computation', default=True)
parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')

###loading data
parser.add_argument('--dataroot', type=str, default='datasets/vangogh2photo/', help='photograph directory root directory')
parser.add_argument('--root2', type=str, default='', help='line drawings dataset root directory')
parser.add_argument('--depthroot', type=str, default='', help='dataset of corresponding ground truth depth maps')
parser.add_argument('--feats2Geom_path', type=str, default='checkpoints/feats2Geom/feats2depth.pth', 
                                help='path to pretrained features to depth map network')

### architecture and optimizers
parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate')
parser.add_argument('--momentum', type=float, default=0.9, help='momentum for optimizer')
parser.add_argument('--decay_epoch', type=int, default=50, help='epoch to start linearly decaying the learning rate to 0')
parser.add_argument('--size', type=int, default=256, help='size of the data crop (squared assumed)')
parser.add_argument('--input_nc', type=int, default=3, help='number of channels of input data')  # DT: changed from 3
parser.add_argument('--output_nc', type=int, default=3, help='number of channels of output data') # DT: changed from 1
parser.add_argument('--geom_nc', type=int, default=3, help='number of channels of geom data')
parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
parser.add_argument('--netD', type=str, default='basic', help='selects model to use for netD')
parser.add_argument('--n_blocks', type=int, default=3, help='number of resnet blocks for generator')
parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')
parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')
parser.add_argument('--disc_sigmoid', type=int, default=0, help='use sigmoid in disc loss')
parser.add_argument('--every_feat', type=int, default=1, help='use transfer features for recog loss')
parser.add_argument('--finetune_netGeom', type=int, default=1, help='make geometry networks trainable')

### loading from checkpoints
parser.add_argument('--load_pretrain', type=str, default='', help='where to load file if wanted')
parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load from if continue_train')

### dataset options
parser.add_argument('--mode', type=str, default='train', help='train, val, test, etc')
parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')
parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')
parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')
parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')
parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')

######## loss weights
parser.add_argument("--cond_cycle", type=float, default=0.1, help="weight of the appearance reconstruction loss")
parser.add_argument("--condGAN", type=float, default=1.0, help="weight of the adversarial style loss")
parser.add_argument("--cond_recog", type=float, default=10.0, help="weight of the semantic loss")
parser.add_argument("--condGeom", type=float, default=10.0, help="weight of the geometry style loss")

### geometry loss options
parser.add_argument('--use_geom', type=int, default=1, help='include the geometry loss')
parser.add_argument('--midas', type=int, default=1, help='use midas depth map')

### semantic loss options
parser.add_argument('--N_patches', type=int, default=1, help='number of patches for clip')
parser.add_argument('--patch_size', type=int, default=128, help='patchsize for clip')
parser.add_argument('--num_classes', type=int, default=55, help='number of classes for inception')
parser.add_argument('--cos_clip', type=int, default=0, help='use cosine similarity for CLIP semantic loss')

### save options
parser.add_argument('--save_epoch_freq', type=int, default=20, help='how often to save the latest model in steps') # DT: changed from 1000
parser.add_argument('--slow', type=int, default=0, help='only frequently save netG_A, netGeom')
parser.add_argument('--log_int', type=int, default=50, help='display frequency for tensorboard')


opt = parser.parse_args()
print(opt)

checkpoints_dir = opt.checkpoints_dir 
name = opt.name

from util.visualizer2 import Visualizer
tensor2im = util.tensor2imv2


visualizer = Visualizer(checkpoints_dir, name, tf_log=True, isTrain=True)
print('------------------- created visualizer -------------------')

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

###### Definition of variables ######
# Networks

netG_A = 0
netG_A = Generator(opt.input_nc, opt.output_nc, opt.n_blocks)
netG_B = Generator(opt.output_nc, opt.input_nc, opt.n_blocks)

if opt.use_geom == 1:

    netGeom = GlobalGenerator2(768, opt.geom_nc, n_downsampling=1, n_UPsampling=3)
    netGeom.load_state_dict(torch.load(opt.feats2Geom_path))
    print("Loading pretrained features to depth network from %s"%opt.feats2Geom_path)
    if opt.finetune_netGeom == 0:
        netGeom.eval()
else:
    opt.finetune_netGeom = 0


D_input_nc_B = opt.output_nc
D_input_nc_A = opt.input_nc

netD_B = networks.define_D(D_input_nc_B, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, use_sigmoid=False)
netD_A = networks.define_D(D_input_nc_A, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, use_sigmoid=False)

device = 'cpu'
if opt.cuda:
    netG_A.cuda()
    netG_B.cuda()
    netD_A.cuda()
    netD_B.cuda()
    if opt.use_geom==1:
        netGeom.cuda()
    device = 'cuda'

### load pretrained inception
net_recog = InceptionV3(opt.num_classes, opt.mode=='test', use_aux=True, pretrain=True, freeze=True, every_feat=opt.every_feat==1)
net_recog.cuda()
net_recog.eval()

import clip
device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
clip.model.convert_weights(clip_model)



#### load in progress weights if continue train or load_pretrain
if opt.continue_train:
    netG_A.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netG_A_%s.pth' % opt.which_epoch)))
    netG_B.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netG_B_%s.pth' % opt.which_epoch)))
    netD_A.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netD_A_%s.pth' % opt.which_epoch)))
    netD_B.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netD_B_%s.pth' % opt.which_epoch)))
    if opt.finetune_netGeom == 1:
        netGeom.load_state_dict(torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'netGeom_%s.pth'% opt.which_epoch)))
    print('----------- loaded %s from '%opt.which_epoch + os.path.join(checkpoints_dir, name) + '---------------------- !!')
elif len(opt.load_pretrain) > 0:
    pretrained_path = opt.load_pretrain
    netG_A.load_state_dict(torch.load(os.path.join(pretrained_path, 'netG_A_%s.pth' % opt.which_epoch)))
    netG_B.load_state_dict(torch.load(os.path.join(pretrained_path, 'netG_B_%s.pth' % opt.which_epoch)))
    netD_A.load_state_dict(torch.load(os.path.join(pretrained_path, 'netD_A_%s.pth' % opt.which_epoch)))
    netD_B.load_state_dict(torch.load(os.path.join(pretrained_path, 'netD_B_%s.pth' % opt.which_epoch)))
    if opt.finetune_netGeom == 1:
        netGeom.load_state_dict(torch.load(os.path.join(pretrained_path, 'netGeom_%s.pth'% opt.which_epoch)))
    print('----------- loaded %s from '%opt.which_epoch + ' ' + pretrained_path + '---------------------- !!')
else:
    netG_A.apply(weights_init_normal)
    netG_B.apply(weights_init_normal)
    netD_A.apply(weights_init_normal)
    netD_B.apply(weights_init_normal)
    

print('----------- loaded networks ---------------------- !!')

# Losses

criterionGAN = networks.GANLoss(use_lsgan=True, target_real_label=1.0, target_fake_label=0.0, reduceme=True).to(device)

criterion_MSE = torch.nn.MSELoss(reduce=True)
criterionCycle = torch.nn.L1Loss()
criterionCycleB = criterionCycle

criterionCLIP = criterion_MSE
if opt.cos_clip == 1:
    criterionCLIP = torch.nn.CosineSimilarity(dim=1, eps=1e-08)

criterionGeom = torch.nn.BCELoss(reduce=True)

############### only use B to A ###########################
optimizer_G_A = torch.optim.Adam(netG_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))
optimizer_G_B = torch.optim.Adam(netG_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))

if (opt.use_geom == 1 and opt.finetune_netGeom==1):
    optimizer_Geom = torch.optim.Adam(netGeom.parameters(), lr=opt.lr, betas=(0.5, 0.999))

optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))
optimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))

lr_scheduler_G_A = torch.optim.lr_scheduler.LambdaLR(optimizer_G_A,
                                                   lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)
lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B,
                                                     lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)

lr_scheduler_G_B = torch.optim.lr_scheduler.LambdaLR(optimizer_G_B,
                                                   lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)
lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A,
                                                     lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step)

# Inputs & targets memory allocation
Tensor = torch.cuda.FloatTensor if opt.cuda else torch.Tensorreal_A

# Dataset loader
transforms_r = [transforms.Resize(int(opt.size * 1.12), Image.BICUBIC),
               transforms.RandomCrop(opt.size),
               transforms.ToTensor()]

train_data = UnpairedDepthDataset(opt.dataroot, opt.root2, opt, transforms_r=transforms_r, 
                mode=opt.mode, midas=opt.midas>0, depthroot=opt.depthroot)

dataloader = DataLoader(train_data, batch_size=opt.batchSize, shuffle=True, num_workers=opt.n_cpu, drop_last=True)


print('---------------- loaded %d images ----------------' % len(train_data))


###### Training ######
for epoch in range(opt.epoch, opt.n_epochs):
    for i, batch in enumerate(dataloader):
        total_steps = epoch*len(dataloader) + i

        img_r  = Variable(batch['r']).cuda()
        img_depth  = Variable(batch['depth']).cuda()

        real_A = img_r
        labels = Variable(batch['label']).cuda()

        real_B = 0
        real_B  = Variable(batch['line']).cuda()

        recover_geom = img_depth
        batch_size = real_A.size()[0]

        condGAN = opt.condGAN
        cond_recog = opt.cond_recog
        cond_cycle = opt.cond_cycle

        #################### Generator ####################


        fake_B = netG_A(real_A) # G_A(A)
        rec_A = netG_B(fake_B)   # G_B(G_A(A))

        fake_A = netG_B(real_B)  # G_B(B)
        rec_B = netG_A(fake_A) # G_A(G_B(B))

        loss_cycle_Geom = 0
        if opt.use_geom == 1:
            geom_input = fake_B
            if geom_input.size()[1] == 1:
                geom_input = geom_input.repeat(1, 3, 1, 1)
            _, geom_input = net_recog(geom_input)

            pred_geom = netGeom(geom_input)
            pred_geom = (pred_geom+1)/2.0 ###[-1, 1] ---> [0, 1]

            loss_cycle_Geom = criterionGeom(pred_geom, recover_geom)

        ########## loss A Reconstruction ##########

        loss_G_A = criterionGAN(netD_A(fake_A), True)

        # GAN loss D_B(G_B(B))
        loss_G_B = 0
        pred_fake_GAN = netD_B(fake_B)
        loss_G_B = criterionGAN(netD_B(fake_B), True)

        # Forward cycle loss || G_B(G_A(A)) - A||
        loss_cycle_A = criterionCycle(rec_A, real_A)
        loss_cycle_B = criterionCycleB(rec_B, real_B)
        # combined loss and calculate gradients

        loss_GAN = loss_G_A + loss_G_B
        loss_RC = loss_cycle_A + loss_cycle_B

        loss_G = cond_cycle*loss_RC + condGAN*loss_GAN
        loss_G += opt.condGeom*loss_cycle_Geom


        ### semantic loss
        loss_recog = 0

        # renormalize mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)
        recog_real = real_A
        recog_real0 = (recog_real[:, 0, :, :].unsqueeze(1) - 0.48145466) / 0.26862954
        recog_real1 = (recog_real[:, 1, :, :].unsqueeze(1) - 0.4578275) / 0.26130258
        recog_real2 = (recog_real[:, 2, :, :].unsqueeze(1) - 0.40821073) / 0.27577711
        recog_real = torch.cat([recog_real0, recog_real1, recog_real2], dim=1)

        line_input = fake_B
        if opt.output_nc == 1:
            line_input_channel0 = (line_input - 0.48145466) / 0.26862954
            line_input_channel1 = (line_input - 0.4578275) / 0.26130258
            line_input_channel2 = (line_input - 0.40821073) / 0.27577711
            line_input = torch.cat([line_input_channel0, line_input_channel1, line_input_channel2], dim=1)

        patches_r = [torch.nn.functional.interpolate(recog_real, size=224)]  #The resize operation on tensor.
        patches_l = [torch.nn.functional.interpolate(line_input, size=224)]

        ## patch based clip loss
        if opt.N_patches > 1:
            patches_r2, patches_l2 = createNRandompatches(recog_real, line_input, opt.N_patches, opt.patch_size)
            patches_r += patches_r2
            patches_l += patches_l2

        loss_recog = 0
        for patchnum in range(len(patches_r)):

            real_patch = patches_r[patchnum]
            line_patch = patches_l[patchnum]

            feats_r = clip_model.encode_image(real_patch).detach()
            feats_line = clip_model.encode_image(line_patch)

            myloss_recog = criterionCLIP(feats_line, feats_r.detach())
            if opt.cos_clip == 1:
                myloss_recog = 1.0 - loss_recog
                myloss_recog = torch.mean(loss_recog)

            patch_factor = (1.0 / float(opt.N_patches))
            if patchnum == 0:
                patch_factor = 1.0
            loss_recog += patch_factor*myloss_recog
        
        loss_G += cond_recog* loss_recog


        optimizer_G_A.zero_grad()
        optimizer_G_B.zero_grad()
        if opt.finetune_netGeom == 1:
            optimizer_Geom.zero_grad()


        loss_G.backward()

        optimizer_G_A.step()
        optimizer_G_B.step()
        if opt.finetune_netGeom == 1:
            optimizer_Geom.step()

        ##########  Discriminator A ##########

        # Fake loss
        pred_fake_A = netD_A(fake_A.detach())
        loss_D_A_fake = criterionGAN(pred_fake_A, False)

        # Real loss

        pred_real_A = netD_A(real_A)
        loss_D_A_real = criterionGAN(pred_real_A, True)

        # Total loss
        loss_D_A = torch.mean(condGAN * (loss_D_A_real + loss_D_A_fake) ) * 0.5

        optimizer_D_A.zero_grad()
        loss_D_A.backward()
        optimizer_D_A.step()

        ##########  Discriminator B ##########

        # Fake loss
        pred_fake_B = netD_B(fake_B.detach())
        loss_D_B_fake = criterionGAN(pred_fake_B, False)

        # Real loss

        pred_real_B = netD_B(real_B)
        loss_D_B_real = criterionGAN(pred_real_B, True)

        # Total loss
        loss_D_B = torch.mean(condGAN * (loss_D_B_real + loss_D_B_fake) ) * 0.5
        optimizer_D_B.zero_grad()
        loss_D_B.backward()
        optimizer_D_B.step()

        # Progress report
        if (i+1)%opt.log_int==0:

            errors = {}
            errors['total_G'] = loss_G.item() if not isinstance(loss_G, (int,float)) else loss_G
            errors['loss_RC'] = torch.mean(loss_RC) if not isinstance(loss_RC, (int,float)) else loss_RC
            errors['loss_cycle_Geom'] = torch.mean(loss_cycle_Geom) if not isinstance(loss_cycle_Geom, (int,float)) else loss_cycle_Geom
            errors['loss_GAN'] = torch.mean(loss_GAN) if not isinstance(loss_GAN, (int,float)) else loss_GANB
            errors['loss_D_B'] = loss_D_B.item() if not isinstance(loss_D_B, (int,float)) else loss_D_B
            errors['loss_D_A'] = loss_D_A.item() if not isinstance(loss_D_A, (int,float)) else loss_D_A
            errors['loss_recog'] = torch.mean(loss_recog) if not isinstance(loss_recog, (int,float)) else loss_recog
           
            visualizer.print_current_errors(epoch, total_steps, errors, 0)
            visualizer.plot_current_errors(errors, total_steps)

            with torch.no_grad():

                input_img = channel2width(real_A)
                if opt.use_geom == 1:
                    pred_geom = channel2width(pred_geom)
                    input_img = torch.cat([input_img, channel2width(recover_geom)], dim=3)
                    
                input_img_fake = channel2width(fake_A)
                rec_A = channel2width(rec_A)

                show_real_B = real_B

                visuals = OrderedDict([('real_A', tensor2im(input_img.data[0])),
                                           ('real_B', tensor2im(show_real_B.data[0])),
                                           ('fake_A', tensor2im(input_img_fake.data[0])),
                                           ('rec_A', tensor2im(rec_A.data[0])),
                                           ('fake_B', tensor2im(fake_B.data[0]))])

                if opt.use_geom == 1:
                    visuals['pred_geom'] = tensor2im(pred_geom.data[0])

                visualizer.display_current_results(visuals, total_steps, epoch)


    # Update learning rates
    lr_scheduler_G_A.step()
    lr_scheduler_G_B.step()
    lr_scheduler_D_A.step()
    lr_scheduler_D_B.step()

    # Save models checkpoints
    # torch.save(netG_A2B.state_dict(), 'output/netG_A2B.pth')
    if (epoch+1) % opt.save_epoch_freq == 0:
        torch.save(netG_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_A_%02d.pth'%(epoch)))
        if opt.finetune_netGeom == 1:
            torch.save(netGeom.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netGeom_%02d.pth'%(epoch)))
        if opt.slow == 0:
            torch.save(netG_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_B_%02d.pth'%(epoch)))
            torch.save(netD_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_A_%02d.pth'%(epoch)))
            torch.save(netD_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_B_%02d.pth'%(epoch)))

    torch.save(netG_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_A_latest.pth'))
    torch.save(netG_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netG_B_latest.pth'))
    torch.save(netD_B.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_B_latest.pth'))
    torch.save(netD_A.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netD_A_latest.pth'))
    if opt.finetune_netGeom == 1:
        torch.save(netGeom.state_dict(), os.path.join(opt.checkpoints_dir, name, 'netGeom_latest.pth'))

###################################