def __init__(self, opt):
"""Initialize the pix2pix class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
self.result_names = [
'real_A', 'mask_M', 'fake_B', 'final_B', 'real_B', 'refer_R',
'fake_B_gradient', 'real_B_gradient'
] ###
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
#self.model_names = ['PCU']
self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake']
#self.loss_names = ['PixWise']
#self.netPCU = networks.PCUnet(opt.input_nc, opt.output_nc, self.gpu_ids)
self.gradient = GetGradient()
if self.isTrain:
self.model_names = ['G', 'D']
else: # during test time, only load G
self.model_names = ['G']
# define networks (both generator and discriminator)
self.netG = networks.generator(opt.input_nc, opt.output_nc,
self.gpu_ids)
if self.isTrain: # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
self.netD = networks.discriminator(
opt.input_nc + opt.output_nc + 1, opt.ndf, opt.n_layers_D,
opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
if self.isTrain:
# define loss functions
self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
self.criterionL1 = InpaintingLoss(
networks.VGG16FeatureExtractor()).to(self.device)
# initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
self.optimizer_G = torch.optim.Adam(filter(
lambda p: p.requires_grad, self.netG.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
#self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
print("Loading the Validation Dataset...")
dataset_val = InitDataset(config.data_root,
img_tf,
mask_tf,
data="val")
print("Loading the Training Dataset...")
dataset_train = InitDataset(config.data_root,
img_tf,
mask_tf,
data="train")
# Loss fucntion
criterion = InpaintingLoss(VGG16FeatureExtractor(),
tv_loss=config.tv_loss).to(device)
# Optimizer
lr = config.finetune_lr if config.finetune else config.initial_lr
if config.optim == "Adam":
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr,
weight_decay=config.weight_decay)
elif config.optim == "SGD":
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
start_iter = 0
sampler=InfiniteSampler(len(dataset_train)),
num_workers=args.n_threads))
print(len(dataset_train))
model = PConvUNet(input_guides=1 if use_depth else 0).to(device)
if args.finetune:
lr = args.lr_finetune
model.freeze_enc_bn = True
else:
lr = args.lr
start_iter = 0
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr)
criterion = InpaintingLoss(VGG16FeatureExtractor()).to(device)
if args.resume:
start_iter = load_ckpt(args.resume, [('model', model)],
[('optimizer', optimizer)])
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('Starting from iter ', start_iter)
for i in tqdm(range(start_iter, args.max_iter)):
model.train()
image, mask, gt = [x.to(device) for x in next(iterator_train)]
if args.mask_root is not None:
guide = image[:, 3:4, :, :]
image = image[:, 0:3, :, :]
dataset = DS(args.root, img_tf)
iterator_train = iter(
data.DataLoader(dataset,
batch_size=args.batch_size,
sampler=InfiniteSampler(len(dataset)),
num_workers=args.n_threads))
print(len(dataset))
model = DFNet().to(device)
lr = args.lr
start_iter = 0
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
criterion = InpaintingLoss().to(device)
if args.resume:
checkpoint = torch.load(args.resume, map_location=device)
model.load_state_dict(checkpoint)
for i in tqdm(range(start_iter, args.max_iter)):
model.train()
img, mask = [x.to(device) for x in next(iterator_train)]
# inpainting
masked = img * mask
# mosaic
"""
loader_val = data.DataLoader(dataset_val,
batch_size=mini_batch,
sampler=RandomSampler(data_source=dataset_val),
num_workers=4)
loaders = {"train": loader_train, "valid": loader_val}
print('model')
# model, criterion, optimizer, scheduler
#model = vgg13().cuda()
#model = resnet18(pretrained=False, progress=True).cuda()
#model = inception_v3().cuda()
model = resnext101_32x8d().cuda()
#criterion = CustomCriterion().cuda()
criterion = InpaintingLoss(VGG16FeatureExtractor()).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[3, 6, 9, 12, 18, 24, 30, 40, 50, 60, 70, 80, 90],
gamma=.5)
print('training')
# model training
runner = dl.SupervisedRunner()
logdir = './logdir'
runner.train(model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
