I want to use cycleGAN on depthmaps for domain adaptation. We are currently training the cycleGAN with two data sets of 2000 images each, but our result are blurred.
import argparse
import itertools
import os
import random
import torch.backends.cudnn as cudnn
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.utils as vutils
from PIL import Image
from tqdm import tqdm
import matplotlib.pylab as plt
from cyclegan_pytorch import DecayLR
from cyclegan_pytorch import Discriminator
from cyclegan_pytorch import Generator
from cyclegan_pytorch import ImageDataset
from cyclegan_pytorch import ReplayBuffer
from cyclegan_pytorch import weights_init
parser = argparse.ArgumentParser(
description="PyTorch implements Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks")
parser.add_argument("--dataroot", type=str, default="./data",
help="path to datasets. (default:./data)")
parser.add_argument("--dataset", type=str, default="depthmaps_white",
help="dataset name. (default:horse2zebra)"
"Option: [apple2orange, summer2winter_yosemite, horse2zebra, monet2photo, "
"cezanne2photo, ukiyoe2photo, vangogh2photo, maps, facades, selfie2anime, "
"iphone2dslr_flower, ae_photos, ]")
parser.add_argument("--epochs", default=200, type=int, metavar="N",
help="number of total epochs to run")
parser.add_argument("--decay_epochs", type=int, default=100,
help="epoch to start linearly decaying the learning rate to 0. (default:100)")
parser.add_argument("-b", "--batch-size", default=1, type=int,
metavar="N",
help="mini-batch size (default: 1), this is the total "
"batch size of all GPUs on the current node when "
"using Data Parallel or Distributed Data Parallel")
parser.add_argument("--lr", type=float, default=0.0002,
help="learning rate. (default:0.0002)")
parser.add_argument("-p", "--print-freq", default=100, type=int,
metavar="N", help="print frequency. (default:100)")
parser.add_argument("--cuda", action="store_true", help="Enables cuda")
parser.add_argument("--netG_A2B", default="", help="path to netG_A2B (to continue training)")
parser.add_argument("--netG_B2A", default="", help="path to netG_B2A (to continue training)")
parser.add_argument("--netD_A", default="", help="path to netD_A (to continue training)")
parser.add_argument("--netD_B", default="", help="path to netD_B (to continue training)")
parser.add_argument("--image-size", type=int, default=128,
help="size of the data crop (squared assumed). (default:128)")
parser.add_argument("--outf", default="./outputs",
help="folder to output images. (default:./outputs).")
parser.add_argument("--manualSeed", type=int,
help="Seed for initializing training. (default:none)")
args = parser.parse_args()
print(args)
try:
os.makedirs(args.outf)
except OSError:
pass
try:
os.makedirs("weights")
except OSError:
pass
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
print("Random Seed: ", args.manualSeed)
random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
cudnn.benchmark = True
if torch.cuda.is_available() and not args.cuda:
print("WARNING: You have a CUDA device, so you should probably run with --cuda")
if args.netG_A2B != "":
netG_A2B.load_state_dict(torch.load(args.netG_A2B))
if args.netG_B2A != "":
netG_B2A.load_state_dict(torch.load(args.netG_B2A))
if args.netD_A != "":
netD_A.load_state_dict(torch.load(args.netD_A))
if args.netD_B != "":
netD_B.load_state_dict(torch.load(args.netD_B))
define loss function (adversarial_loss) and optimizer
Hello everybody,
I want to use cycleGAN on depthmaps for domain adaptation. We are currently training the cycleGAN with two data sets of 2000 images each, but our result are blurred.
Trainining configuration:
image_size = 128 input_channels = 1 output_channels = 1 norm = instance_norm lr = 0.0002 beta_1 = 0.5 cycle_loss_weight = 10 batch_size = 1 9 res_blocks
Model
import torch import torch.nn as nn import torch.nn.functional as F
class Discriminator(nn.Module): def init(self): super(Discriminator, self).init()
class Generator(nn.Module): def init(self): super(Generator, self).init() self.main = nn.Sequential(
Initial convolution block
class ResidualBlock(nn.Module): def init(self, in_channels): super(ResidualBlock, self).init()
Train
import argparse import itertools import os import random
import torch.backends.cudnn as cudnn import torch.utils.data import torchvision.transforms as transforms import torchvision.utils as vutils from PIL import Image from tqdm import tqdm import matplotlib.pylab as plt
from cyclegan_pytorch import DecayLR from cyclegan_pytorch import Discriminator from cyclegan_pytorch import Generator from cyclegan_pytorch import ImageDataset from cyclegan_pytorch import ReplayBuffer from cyclegan_pytorch import weights_init
parser = argparse.ArgumentParser( description="PyTorch implements
Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks
") parser.add_argument("--dataroot", type=str, default="./data", help="path to datasets. (default:./data)") parser.add_argument("--dataset", type=str, default="depthmaps_white", help="dataset name. (default:horse2zebra
)" "Option: [apple2orange, summer2winter_yosemite, horse2zebra, monet2photo, " "cezanne2photo, ukiyoe2photo, vangogh2photo, maps, facades, selfie2anime, " "iphone2dslr_flower, ae_photos, ]") parser.add_argument("--epochs", default=200, type=int, metavar="N", help="number of total epochs to run") parser.add_argument("--decay_epochs", type=int, default=100, help="epoch to start linearly decaying the learning rate to 0. (default:100)") parser.add_argument("-b", "--batch-size", default=1, type=int, metavar="N", help="mini-batch size (default: 1), this is the total " "batch size of all GPUs on the current node when " "using Data Parallel or Distributed Data Parallel") parser.add_argument("--lr", type=float, default=0.0002, help="learning rate. (default:0.0002)") parser.add_argument("-p", "--print-freq", default=100, type=int, metavar="N", help="print frequency. (default:100)") parser.add_argument("--cuda", action="store_true", help="Enables cuda") parser.add_argument("--netG_A2B", default="", help="path to netG_A2B (to continue training)") parser.add_argument("--netG_B2A", default="", help="path to netG_B2A (to continue training)") parser.add_argument("--netD_A", default="", help="path to netD_A (to continue training)") parser.add_argument("--netD_B", default="", help="path to netD_B (to continue training)") parser.add_argument("--image-size", type=int, default=128, help="size of the data crop (squared assumed). (default:128)") parser.add_argument("--outf", default="./outputs", help="folder to output images. (default:./outputs
).") parser.add_argument("--manualSeed", type=int, help="Seed for initializing training. (default:none)")args = parser.parse_args() print(args)
try: os.makedirs(args.outf) except OSError: pass
try: os.makedirs("weights") except OSError: pass
if args.manualSeed is None: args.manualSeed = random.randint(1, 10000) print("Random Seed: ", args.manualSeed) random.seed(args.manualSeed) torch.manual_seed(args.manualSeed)
cudnn.benchmark = True
if torch.cuda.is_available() and not args.cuda: print("WARNING: You have a CUDA device, so you should probably run with --cuda")
Dataset
dataset = ImageDataset(root=os.path.join(args.dataroot, args.dataset), transform=transforms.Compose([ transforms.Resize(int(args.image_size * 1.12), Image.BICUBIC), transforms.RandomCrop(args.image_size),
transforms.RandomHorizontalFlip(),
dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True)
try: os.makedirs(os.path.join(args.outf, args.dataset, "Sim")) os.makedirs(os.path.join(args.outf, args.dataset, "Real")) except OSError: pass
try: os.makedirs(os.path.join("weights", args.dataset)) except OSError: pass
device = torch.device("cuda:0" if args.cuda else "cpu")
create model
netG_A2B = Generator().to(device) netG_B2A = Generator().to(device) netD_A = Discriminator().to(device) netD_B = Discriminator().to(device)
netG_A2B.apply(weights_init) netG_B2A.apply(weights_init) netD_A.apply(weights_init) netD_B.apply(weights_init)
if args.netG_A2B != "": netG_A2B.load_state_dict(torch.load(args.netG_A2B)) if args.netG_B2A != "": netG_B2A.load_state_dict(torch.load(args.netG_B2A)) if args.netD_A != "": netD_A.load_state_dict(torch.load(args.netD_A)) if args.netD_B != "": netD_B.load_state_dict(torch.load(args.netD_B))
define loss function (adversarial_loss) and optimizer
cycle_loss = torch.nn.L1Loss().to(device) identity_loss = torch.nn.L1Loss().to(device) adversarial_loss = torch.nn.MSELoss().to(device)
Optimizers
optimizer_G = torch.optim.Adam(itertools.chain(netG_A2B.parameters(), netG_B2A.parameters()), lr=args.lr, betas=(0.5, 0.999)) optimizer_D_A = torch.optim.Adam(netD_A.parameters(), lr=args.lr, betas=(0.5, 0.999)) optimizer_D_B = torch.optim.Adam(netD_B.parameters(), lr=args.lr, betas=(0.5, 0.999))
lr_lambda = DecayLR(args.epochs, 0, args.decay_epochs).step lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda=lr_lambda) lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(optimizer_D_A, lr_lambda=lr_lambda) lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(optimizer_D_B, lr_lambda=lr_lambda)
g_losses = [] d_losses = []
identity_losses = [] gan_losses = [] cycle_losses = []
fake_A_buffer = ReplayBuffer() fake_B_buffer = ReplayBuffer()
for epoch in range(0, args.epochs): progress_bar = tqdm(enumerate(dataloader), total=len(dataloader)) for i, data in progress_bar:
get batch size data
save last check pointing
torch.save(netG_A2B.state_dict(), f"weights/{args.dataset}/netG_A2B.pth") torch.save(netG_B2A.state_dict(), f"weights/{args.dataset}/netG_B2A.pth") torch.save(netD_A.state_dict(), f"weights/{args.dataset}/netD_A.pth") torch.save(netD_B.state_dict(), f"weights/{args.dataset}/netD_B.pth")