def main():
# instantiate model and initialize weights
model = ENet()
networks.print_network(model)
networks.init_weights(model, init_type='normal')
model.init_convFilter(trainable=srm_trainable)
if args.cuda:
model.cuda()
torch.save({
'epoch': 0,
'state_dict': model.state_dict()
}, '{}/checkpoint_{}.pth'.format(LOG_DIR, 0))
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
L1_criterion = nn.L1Loss(reduction='sum').cuda()
if not srm_trainable:
params = []
for name, param in model.named_parameters():
if name.find('convFilter1') == -1:
params += [param]
optimizer = create_optimizer(params, args.lr)
else:
optimizer = create_optimizer(model.parameters(), args.lr)
for epoch in range(1, args.epochs + 1):
# update the optimizer learning rate
adjust_learning_rate(optimizer, epoch)
train(train_loader, model, optimizer, criterion, L1_criterion, epoch)
def main():
# instantiate model and initialize weights
model = AutoNet(input_nc=args.input_nc, ndf=6, nonlinear='relu')
networks.print_network(model)
if args.cuda:
model.cuda()
print('using pretrained model')
checkpoint = torch.load(project_root + args.log_dir + '/checkpoint_60.pth')
model.load_state_dict(checkpoint['state_dict'])
args.lr = args.lr * 0.001
itr_start = 1
nature = test_nature(val_loader, model)
print(nature)
threshold = max(2, nature * 2)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = create_optimizer(model, args.lr)
new_samples = []
pn_num = []
nature_error_itr_global = []
for itr in np.arange(itr_start, 5):
args.dataroot = dst_dir
tmp = construct_negative_samples(model, new_samples, itr)
pn_num.append(tmp)
train_loader = myDataset.DataLoaderHalf(
myDataset.MyDataset(
args,
transforms.Compose([
transforms.Resize((256, 256), Image.BICUBIC),
transforms.ToTensor(), normalize
])),
batch_size=args.batch_size,
shuffle=True,
half_constraint=True,
sampler_type='RandomBalancedSampler',
**kwargs)
print('The number of train data:{}'.format(len(train_loader.dataset)))
args.epochs = 15 # after the new negative samples construct, the learning rate is constant, and epoch = 15
train_multi(train_loader, optimizer, model, criterion, val_loader, itr,
nature_error_itr_global)
print(pn_num)
model_selection(nature_error_itr_global, threshold)
def __init__(self, p):
super(TestModel, self).__init__(p)
assert (not p.isTrain)
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
opt.which_model_netG, opt.norm,
not opt.no_dropout, opt.init_type,
self.gpu_ids)
which_epoch = p.which_epoch
self.load_model(self.netG, 'G', which_epoch)
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
print('-----------------------------------------------')
def initialize(self, opt, train_mode=True):
# Model transforms from A --> B and uses Adv as the
# adversarial example.
BaseModel.initialize(self, opt)
self.train_mode = train_mode
# define tensors
self.input_B = self.Tensor(opt['batchSize'], opt['input_nc'],
opt['B_height'], opt['B_width'])
self.input_A = self.Tensor(opt['batchSize'], opt['output_nc'],
opt['A_height'], opt['A_width'])
# load/define networks
self.netG = networks.define_G(opt['input_nc'], opt['output_nc'],
opt['ngf'], opt['norm'], self.gpu_ids)
if self.train_mode:
use_sigmoid = opt['no_lsgan']
self.netD = networks.define_D(opt['input_nc'] + opt['output_nc'],
opt['ndf'], opt['which_model_netD'],
opt['n_layers_D'], use_sigmoid,
self.gpu_ids)
if self.train_mode:
# self.fake_AB_pool = ImagePool(opt['pool_size'])
self.old_lr = opt['lr']
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt['no_lsgan'],
tensor=self.Tensor)
self.content_loss = torch.nn.MSELoss()
# initialize optimizers
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))
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
networks.print_network(self.netD)
print('-----------------------------------------------')
def __init__(self, opt, ignore_noise=False, testing=False):
self.ignore_noise = ignore_noise
##### model options
self.old_lr = opt.lr
opt.use_sigmoid = opt.no_lsgan
self.opt = opt
##### define all networks we need here
self.netG_A_B = networks.define_stochastic_G(nlatent=opt.nlatent, input_nc=opt.input_nc,
output_nc=opt.output_nc, ngf=opt.ngf,
which_model_netG=opt.which_model_netG,
norm=opt.norm, use_dropout=opt.use_dropout,
gpu_ids=opt.gpu_ids)
self.netG_B_A = networks.define_G(input_nc=opt.output_nc,
output_nc=opt.input_nc, ngf=opt.ngf,
which_model_netG=opt.which_model_netG,
norm=opt.norm, use_dropout=opt.use_dropout,
gpu_ids=opt.gpu_ids)
self.netD_A = networks.define_D_A(input_nc=opt.input_nc,
ndf=32, which_model_netD=opt.which_model_netD,
norm=opt.norm, use_sigmoid=opt.use_sigmoid, gpu_ids=opt.gpu_ids)
self.netD_B = networks.define_D_B(input_nc=opt.output_nc,
ndf=opt.ndf, which_model_netD=opt.which_model_netD,
norm=opt.norm, use_sigmoid=opt.use_sigmoid, gpu_ids=opt.gpu_ids)
##### define all optimizers here
self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A_B.parameters(),
self.netG_B_A.parameters()),
lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(),
self.netD_B.parameters()),
lr=opt.lr/5., betas=(opt.beta1, 0.999))
self.criterionGAN = functools.partial(criterion_GAN, use_sigmoid=opt.use_sigmoid)
self.criterionCycle = F.l1_loss
if not testing:
with open("%s/nets.txt" % opt.expr_dir, 'w') as nets_f:
networks.print_network(self.netG_A_B, nets_f)
networks.print_network(self.netG_B_A, nets_f)
networks.print_network(self.netD_A, nets_f)
networks.print_network(self.netD_B, nets_f)
def initialize(self, args):
BaseModel.initialize(self, args)
self.nb = args['batch_size']
sizeH, sizeW = args['fineSizeH'], args['fineSizeW']
self.input_A = self.Tensor(self.nb, args['input_nc'], sizeH, sizeW)
self.input_B = self.Tensor(self.nb, args['input_nc'], sizeH, sizeW)
self.input_A_label = torch.cuda.LongTensor(self.nb, args['input_nc'],
sizeH, sizeW)
self.netG = networks.netG().cuda(device_id=args['device_ids'][0])
self.netD = define_D(
args['net_D']).cuda(device_id=args['device_ids'][0])
self.deeplabPart1 = networks.DeeplabPool1().cuda(
device_id=args['device_ids'][0])
self.deeplabPart2 = networks.DeeplabPool12Pool5().cuda(
device_id=args['device_ids'][0])
self.deeplabPart3 = networks.DeeplabPool52Fc8_interp().cuda(
device_id=args['device_ids'][0])
# define loss functions
self.criterionCE = torch.nn.CrossEntropyLoss(size_average=False)
self.criterionAdv = networks.Advloss(use_lsgan=args['use_lsgan'],
tensor=self.Tensor)
if not args['resume']:
#initialize networks
self.netG.apply(weights_init)
self.netD.apply(weights_init)
pretrained_dict = torch.load(args['weigths_pool'] + '/' +
args['pretrain_model'])
self.deeplabPart1.weights_init(pretrained_dict=pretrained_dict)
self.deeplabPart2.weights_init(pretrained_dict=pretrained_dict)
self.deeplabPart3.weights_init(pretrained_dict=pretrained_dict)
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=args['lr_gan'],
betas=(args['beta1'], 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=args['lr_gan'],
betas=(args['beta1'], 0.999))
ignored_params = list(map(id, self.deeplabPart3.fc8_1.parameters()))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_2.parameters())))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_3.parameters())))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_4.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
self.deeplabPart3.parameters())
base_params = base_params + filter(lambda p: True,
self.deeplabPart1.parameters())
base_params = base_params + filter(lambda p: True,
self.deeplabPart2.parameters())
self.optimizer_P = torch.optim.SGD([
{
'params': base_params
},
{
'params': get_parameters(self.deeplabPart3.fc8_1, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_2, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_3, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_4, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_1, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_2, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_3, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_4, 'bias'),
'lr': args['l_rate'] * 20
},
],
lr=args['l_rate'],
momentum=0.9,
weight_decay=5e-4)
#netG_params = filter(lambda p: True, self.netG.parameters())
self.optimizer_R = torch.optim.SGD(
[
{
'params': base_params
},
#{'params': netG_params, 'lr': args['l_rate'] * 100},
{
'params': get_parameters(self.deeplabPart3.fc8_1,
'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_2,
'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_3,
'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_4,
'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_1, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_2, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_3, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_4, 'bias'),
'lr': args['l_rate'] * 20
},
],
lr=args['l_rate'],
momentum=0.9,
weight_decay=5e-4)
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
networks.print_network(self.netD)
networks.print_network(self.deeplabPart1)
networks.print_network(self.deeplabPart2)
networks.print_network(self.deeplabPart3)
print('-----------------------------------------------')
netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, [0])
netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0])
print('loading done')
criterionGAN = GANLoss()
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
# setup optimizer
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
print('---------- Networks initialized -------------')
print_network(netG)
print_network(netD)
print('-----------------------------------------------')
real_a = torch.FloatTensor(opt.batchSize, opt.input_nc, 256, 256)
real_b = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256)
if opt.cuda:
netD = netD.cuda()
netG = netG.cuda()
criterionGAN = criterionGAN.cuda()
criterionL1 = criterionL1.cuda()
criterionMSE = criterionMSE.cuda()
real_a = real_a.cuda()
real_b = real_b.cuda()
def initialize(self, args):
BaseModel.initialize(self, args)
self.if_adv_train = args['if_adv_train']
self.Iter = 0
self.interval_g2 = args['interval_g2']
self.interval_d2 = args['interval_d2']
self.nb = args['batch_size']
sizeH, sizeW = args['fineSizeH'], args['fineSizeW']
self.tImageA = self.Tensor(self.nb, args['input_nc'], sizeH, sizeW)
self.tImageB = self.Tensor(self.nb, args['input_nc'], sizeH, sizeW)
self.tLabelA = torch.cuda.LongTensor(self.nb, 1, sizeH, sizeW)
self.tOnehotLabelA = self.Tensor(self.nb, args['label_nums'], sizeH,
sizeW)
self.loss_G = Variable()
self.loss_D = Variable()
self.netG1 = networks.netG().cuda(device_id=args['device_ids'][0])
self.netD1 = define_D(args['net_d1'],
512).cuda(device_id=args['device_ids'][0])
self.netD2 = define_D(
args['net_d2'],
args['label_nums']).cuda(device_id=args['device_ids'][0])
self.deeplabPart1 = networks.DeeplabPool1().cuda(
device_id=args['device_ids'][0])
self.deeplabPart2 = networks.DeeplabPool12Pool5().cuda(
device_id=args['device_ids'][0])
self.deeplabPart3 = networks.DeeplabPool52Fc8_interp(
output_nc=args['label_nums']).cuda(device_id=args['device_ids'][0])
# define loss functions
self.criterionCE = torch.nn.CrossEntropyLoss(size_average=False)
self.criterionAdv = networks.Advloss(use_lsgan=args['use_lsgan'],
tensor=self.Tensor)
if not args['resume']:
#initialize networks
self.netG1.apply(weights_init)
self.netD1.apply(weights_init)
self.netD2.apply(weights_init)
pretrained_dict = torch.load(args['weigths_pool'] + '/' +
args['pretrain_model'])
self.deeplabPart1.weights_init(pretrained_dict=pretrained_dict)
self.deeplabPart2.weights_init(pretrained_dict=pretrained_dict)
self.deeplabPart3.weights_init(pretrained_dict=pretrained_dict)
# initialize optimizers
self.optimizer_G1 = torch.optim.Adam(self.netG1.parameters(),
lr=args['lr_g1'],
betas=(args['beta1'], 0.999))
self.optimizer_D1 = torch.optim.Adam(self.netD1.parameters(),
lr=args['lr_g1'],
betas=(args['beta1'], 0.999))
self.optimizer_G2 = torch.optim.Adam(
[{
'params': self.deeplabPart1.parameters()
}, {
'params': self.deeplabPart2.parameters()
}, {
'params': self.deeplabPart3.parameters()
}],
lr=args['lr_g2'],
betas=(args['beta1'], 0.999))
self.optimizer_D2 = torch.optim.Adam(self.netD2.parameters(),
lr=args['lr_g2'],
betas=(args['beta1'], 0.999))
ignored_params = list(map(id, self.deeplabPart3.fc8_1.parameters()))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_2.parameters())))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_3.parameters())))
ignored_params.extend(
list(map(id, self.deeplabPart3.fc8_4.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
self.deeplabPart3.parameters())
base_params = base_params + filter(lambda p: True,
self.deeplabPart1.parameters())
base_params = base_params + filter(lambda p: True,
self.deeplabPart2.parameters())
deeplab_params = [
{
'params': base_params
},
{
'params': get_parameters(self.deeplabPart3.fc8_1, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_2, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_3, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_4, 'weight'),
'lr': args['l_rate'] * 10
},
{
'params': get_parameters(self.deeplabPart3.fc8_1, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_2, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_3, 'bias'),
'lr': args['l_rate'] * 20
},
{
'params': get_parameters(self.deeplabPart3.fc8_4, 'bias'),
'lr': args['l_rate'] * 20
},
]
self.optimizer_P = torch.optim.SGD(deeplab_params,
lr=args['l_rate'],
momentum=0.9,
weight_decay=5e-4)
self.optimizer_R = torch.optim.SGD(deeplab_params,
lr=args['l_rate'],
momentum=0.9,
weight_decay=5e-4)
print('---------- Networks initialized -------------')
networks.print_network(self.netG1)
networks.print_network(self.netD1)
networks.print_network(self.netD2)
networks.print_network(self.deeplabPart1)
networks.print_network(self.deeplabPart2)
networks.print_network(self.deeplabPart3)
print('-----------------------------------------------')
#netD_Edge = define_D_Edge(hogehoge)
criterionGAN = GANLoss()
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
# setup optimizer
optimizerG = optim.Adadelta(netG.parameters(), lr=opt.lr)
optimizerD_Global = optim.Adadelta(netD_Global.parameters(), lr=opt.lr)
optimizerD_Local = optim.Adadelta(netD_Local.parameters(), lr=opt.lr)
optimizerD_Edge = optim.Adadelta(netD_Edge.parameters(), lr=opt.lr)
print('---------- Networks initialized -------------')
print_network(netG)
print_network(netD_Global)
print_network(netD_Local)
print_network(netD_Edge)
print('-----------------------------------------------')
real_a0_image = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256)
real_a_image = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256)
#基本的なテンソルに使うShape
#tensor_shape = torch.tensor(opt.batchSize,3,image_size,image_size)
tensor_shape = torch.empty([opt.batchSize,3,image_size,image_size])
if opt.cuda:
#netD = netD.cuda()
netD_Global = netD_Global.cuda()
netD_Local = netD_Local.cuda()
def main(name_exp, segloss=False, cuda=True, finetune=False):
# Training settings
parser = argparse.ArgumentParser(description='pix2pix-PyTorch-implementation')
parser.add_argument('--batchSize', type=int, default=8, help='training batch size')
parser.add_argument('--testBatchSize', type=int, default=8, help='testing batch size')
parser.add_argument('--nEpochs', type=int, default=100, help='number of epochs to train for')
parser.add_argument('--input_nc', type=int, default=3, help='input image channels')
parser.add_argument('--output_nc', type=int, default=3, help='output image channels')
parser.add_argument('--ngf', type=int, default=64, help='generator filt+ers in first conv layer')
parser.add_argument('--ndf', type=int, default=64, help='discriminator filters in first conv layer')
parser.add_argument('--lr', type=float, default=0.0002, help='Learning Rate. Default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--threads', type=int, default=8, help='number of threads for data loader to use')
parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
parser.add_argument('--lamb', type=int, default=10, help='weight on L1 term in objective')
opt = parser.parse_args()
cudnn.benchmark = True
def val():
net_current = "path_exp/checkpoint/DFS/{}/netG_model_current.pth".format(name_exp)
netVal = torch.load(net_current)
netVal.eval()
SEG_NET.eval()
features.eval()
with torch.no_grad():
total_mse = 0
total_mse2 = 0
avg_psnr_depth = 0
avg_psnr_dehaze = 0
avg_ssim_depth = 0
avg_ssim_dehaze = 0
for batch in validation_data_loader:
input, target, depth = Variable(batch[0]), Variable(batch[1]), Variable(batch[2])
if cuda == True:
input = input.cuda()
target = target.cuda()
depth = depth.cuda()
dehaze = netVal(input)
prediction = SEG_NET(dehaze)
avg_ssim_dehaze += pytorch_ssim.ssim(dehaze, target).item()
mse = criterionMSE(prediction, depth)
total_mse += mse.item()
avg_psnr_depth += 10 * log10(1 / mse.item())
mse2 = criterionMSE(dehaze, target)
total_mse2 += mse2.item()
avg_psnr_dehaze += 10 * log10(1 / mse2.item())
avg_ssim_depth += pytorch_ssim.ssim(prediction, depth).item()
visual_ret_val = OrderedDict()
visual_ret_val['Haze'] = input
visual_ret_val['Seg estimate'] = prediction
visual_ret_val['Dehaze '] = dehaze
visual_ret_val['GT dehaze'] = target
visual_ret_val['GT Seg '] = depth
visualizer.display_current_results(visual_ret_val, epoch, True)
print("===> Validation")
#f.write("===> Testing: \r\n")
print("===> PSNR seg: {:.4f} ".format(avg_psnr_depth / len(validation_data_loader)))
#f.write("===> PSNR depth: {:.4f} \r\n".format(avg_psnr_depth / len(validation_data_loader)))
print("===> Mse seg: {:.4f} ".format(total_mse / len(validation_data_loader)))
#f.write("===> Mse depth: {:.4f} \r\n".format(total_mse / len(validation_data_loader)))
print("===> SSIM seg: {:.4f} ".format(avg_ssim_depth / len(validation_data_loader)))
#f.write("===> SSIM depth: {:.4f} \r\n".format(avg_ssim_depth / len(validation_data_loader)))
return total_mse / len(validation_data_loader)
def testing():
path = "path_exp/checkpoint/DFS/{}/netG_model_best.pth".format(name_exp)
net = torch.load(path)
net.eval()
SEG_NET.eval()
features.eval()
with torch.no_grad():
total_mse = 0
total_mse2 = 0
avg_psnr_depth = 0
avg_psnr_dehaze = 0
avg_ssim_depth = 0
avg_ssim_dehaze = 0
for batch in testing_data_loader:
input, target, depth = Variable(batch[0]), Variable(batch[1]), Variable(batch[2])
if cuda == True:
input = input.cuda()
target = target.cuda()
depth = depth.cuda()
dehaze = net(input)
prediction = SEG_NET(dehaze)
avg_ssim_dehaze += pytorch_ssim.ssim(dehaze, target).item()
mse = criterionMSE(prediction, depth)
total_mse += mse.item()
avg_psnr_depth += 10 * log10(1 / mse.item())
mse2 = criterionMSE(dehaze, target)
total_mse2 += mse2.item()
avg_psnr_dehaze += 10 * log10(1 / mse2.item())
avg_ssim_depth += pytorch_ssim.ssim(prediction, depth).item()
print("===> Testing")
print("===> PSNR seg: {:.4f} ".format(avg_psnr_depth / len(testing_data_loader)))
print("===> Mse seg: {:.4f} ".format(total_mse / len(testing_data_loader)))
print("===> SSIM seg: {:.4f} ".format(avg_ssim_depth / len(testing_data_loader)))
print("===> PSNR dehaze: {:.4f} ".format(avg_psnr_dehaze / len(testing_data_loader)))
print("===> SSIM dehaze: {:.4f} ".format(avg_ssim_dehaze / len(testing_data_loader)))
def checkpoint():
if not os.path.exists("checkpoint"):
os.mkdir("checkpoint")
if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)):
os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp))
net_g_model_out_path = "path_exp/checkpoint/DFS/{}/netG_model_best.pth".format(name_exp)
net_d_model_out_path = "path_exp/checkpoint/DFS/{}/netD_model_best.pth".format(name_exp)
torch.save(netG, net_g_model_out_path)
torch.save(netD, net_d_model_out_path)
def checkpoint_current():
if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)):
os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp))
net_g_model_out_path = "path_exp/checkpoint/DFS/{}/netG_model_current.pth".format(name_exp)
torch.save(netG, net_g_model_out_path)
def checkpoint_seg():
if not os.path.exists(os.path.join("path_exp/checkpoint/DFS", name_exp)):
os.mkdir(os.path.join("path_exp/checkpoint/DFS", name_exp))
net_g_model_out_path = "path_exp/checkpoint/DFS/{}/seg_net.pth".format(name_exp)
torch.save(SEG_NET, net_g_model_out_path)
torch.manual_seed(opt.seed)
if cuda==True:
torch.cuda.manual_seed(opt.seed)
print(" ")
print(name_exp)
print(" ")
print('===> Loading datasets')
train_set = get_training_set('path_exp/cityscape/HAZE')
val_set = get_val_set('path_exp/cityscape/HAZE')
test_set = get_test_set('path_exp/cityscape/HAZE')
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
validation_data_loader = DataLoader(dataset=val_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
testing_data_loader= DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model')
netG = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, [0])
netD = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'batch', False, [0])
criterionGAN = GANLoss()
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
# setup optimizer
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
print('---------- Networks initialized -------------')
print_network(netG)
print_network(netD)
print('-----------------------------------------------')
real_a = torch.FloatTensor(opt.batchSize, opt.input_nc, 256, 256)
real_b = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256)
real_c = torch.FloatTensor(opt.batchSize, opt.output_nc, 256, 256)
if cuda==True:
netD = netD.cuda()
netG = netG.cuda()
criterionGAN = criterionGAN.cuda()
criterionL1 = criterionL1.cuda()
criterionMSE = criterionMSE.cuda()
real_a = real_a.cuda()
real_b = real_b.cuda()
real_c=real_c.cuda()
real_a = Variable(real_a)
real_b = Variable(real_b)
real_c = Variable(real_c)
SEG_NET = torch.load("path_exp/SEG_NET.pth")
optimizerSeg = optim.Adam(SEG_NET.parameters(), lr=opt.lr/10, betas=(opt.beta1, 0.999))
features = Vgg16()
if cuda==True:
SEG_NET.cuda()
features.cuda()
bon =100000000
for epoch in range(opt.nEpochs):
features.eval()
if finetune== True and epoch>50:
SEG_NET.train()
else:
SEG_NET.eval()
loss_epoch_gen=0
loss_epoch_dis=0
total_segloss=0
loss_seg=0
i=0
for iteration, batch in enumerate(training_data_loader, 1):
netG.train()
i=i+1
# forward
real_a_cpu, real_b_cpu, real_c_cpu = batch[0], batch[1], batch[2]
with torch.no_grad():
real_a = real_a.resize_(real_a_cpu.size()).copy_(real_a_cpu)
with torch.no_grad():
real_b = real_b.resize_(real_b_cpu.size()).copy_(real_b_cpu)
with torch.no_grad():
real_c = real_c.resize_(real_c_cpu.size()).copy_(real_c_cpu)
fake_b = netG(real_a)
############################
# (1) Update D network: maximize log(D(x,y)) + log(1 - D(x,G(x)))
###########################
optimizerD.zero_grad()
# train with fake
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = netD.forward(fake_ab.detach())
loss_d_fake = criterionGAN(pred_fake, False)
# train with real
real_ab = torch.cat((real_a, real_b), 1)
pred_real = netD.forward(real_ab)
loss_d_real = criterionGAN(pred_real, True)
# Combined loss
loss_d = (loss_d_fake + loss_d_real) * 0.5
loss_d.backward()
optimizerD.step()
############################
# (2) Update G network: maximize log(D(x,G(x))) + L1(y,G(x))
##########################
optimizerG.zero_grad()
# First, G(A) should fake the discriminator
fake_ab = torch.cat((real_a, fake_b), 1)
pred_fake = netD.forward(fake_ab)
loss_g_gan = criterionGAN(pred_fake, True)
# Second, G(A) = B
loss_g_l1 = criterionL1(fake_b, real_b) * opt.lamb
features_y = features(fake_b)
features_x = features(real_b)
loss_content = criterionMSE(features_y[1], features_x[1])*10
if segloss == True:
fake_seg = SEG_NET(fake_b)
loss_seg = criterionMSE(fake_seg, real_c) * 10
total_segloss += loss_seg.item()
features_y = features(fake_seg)
features_x = features(real_c)
ssim_seg = criterionMSE(features_y[1], features_x[1]) * 10
loss_g = loss_g_gan + loss_g_l1 + loss_content + loss_seg
else:
loss_g = loss_g_gan + loss_g_l1+loss_content
loss_epoch_gen+=loss_g.item()
loss_epoch_dis+=loss_d.item()
if finetune== True and epoch>50:
loss_g.backward(retain_graph=True)
optimizerG.step()
loss_seg=loss_seg
loss_seg.backward()
optimizerSeg.zero_grad()
optimizerSeg.step()
else:
loss_g.backward()
optimizerG.step()
errors_ret = OrderedDict()
errors_ret['Total_G'] = float(loss_g)
errors_ret['Content'] = float(loss_content)
errors_ret['GAN'] = float(loss_g_gan)
errors_ret['L1'] = float(loss_g_l1)
errors_ret['D'] = float(loss_d)
if i % 10 == 0: # print training losses and save logging information to the disk
if i > 0:
visualizer.plot_current_losses(epoch, i/(len(training_data_loader)*opt.batchSize), errors_ret)
print("===> Epoch[{}]: Loss_D: {:.4f} Loss_G: {:.4f} Loss Seg: {:.4f} ".format(epoch, loss_epoch_dis,loss_epoch_gen, total_segloss))
checkpoint_current()
MSE=val()
if MSE < bon:
bon = MSE
checkpoint()
checkpoint_seg()
print("BEST EPOCH SAVED")
testing()
def main():
# instantiate model and initialize weights
model = ENet()
networks.print_network(model)
networks.init_weights(model, init_type='normal')
model.init_convFilter(trainable=srm_trainable)
if args.cuda:
model.cuda()
print('using pretrained model')
checkpoint = torch.load(project_root + args.log_dir +
'/checkpoint_300.pth')
model.load_state_dict(checkpoint['state_dict'])
args.lr = args.lr * 0.001
threshold = THRESHOLD_MAX
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
L1_criterion = nn.L1Loss(reduction='sum').cuda()
if not srm_trainable:
params = []
for name, param in model.named_parameters():
if name.find('convFilter1') == -1:
params += [param]
optimizer = create_optimizer(params, args.lr)
else:
optimizer = create_optimizer(model.parameters(), args.lr)
nature_error_itr_global = []
for itr in np.arange(1, 11):
args.dataroot = dst_dir
nature_error_itr_local = []
# adding negative samples into the original training dataset
construct_negative_samples(itr)
train_loader = myDataset.DataLoaderHalf(
myDataset.MyDataset(
args,
transforms.Compose([
transforms.RandomCrop(233),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
])),
batch_size=args.batch_size,
shuffle=True,
half_constraint=True,
sampler_type='RandomBalancedSampler',
**kwargs)
print('The number of train data:{}'.format(len(train_loader.dataset)))
args.epochs = 15
train_multi(train_loader, optimizer, model, criterion, L1_criterion, val_loader, itr, \
nature_error_itr_local, nature_error_itr_global)
# start from itr = 1
if len(nature_error_itr_local) > 0:
adv_model_num, adv_model_idx = adv_model_selection(
nature_error_itr_local, threshold, itr)
if adv_model_num < 1:
break
print(nature_error_itr_global)
print(len(nature_error_itr_global) / (args.epochs - args.epochs // 2))
final_model_selection(nature_error_itr_global, threshold)
def print_net(self):
networks.print_network(self.net)
# Set up loss
criterionGAN = GANLoss()
criterionL1 = nn.L1Loss()
criterionMSE = nn.MSELoss()
criterionSLL = SLLoss()
# Set up optimizer
optimizerG = optim.Adam(list(G.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
optimizerD = optim.Adam(list(D.parameters()),
lr=opt.lr,
betas=(opt.beta1, 0.999))
print_network(G)
print_network(D)
# Train
real_a = torch.FloatTensor(opt.batchSize, opt.input_channels, 256, 256)
real_b = torch.FloatTensor(opt.batchSize, opt.output_channels, 256, 256)
coords = torch.FloatTensor(opt.batchSize, 4)
if opt.cuda:
D = D.cuda()
G = G.cuda()
criterionGAN = criterionGAN.cuda()
criterionSLL = criterionSLL.cuda()
criterionL1 = criterionL1.cuda()
criterionMSE = criterionMSE.cuda()
real_a = real_a.cuda()
def __init__(
self, name="experiment", phase="train", which_epoch="latest",
batch_size=1, image_size=128, map_nc=1, input_nc=3, output_nc=3,
num_downs=7, ngf=64, ndf=64, norm_layer="batch", pool_size=50,
lr=0.0002, beta1=0.5, lambda_D=0.5, lambda_MSE=10,
lambda_P=5.0, use_dropout=True, gpu_ids=[], n_layers=3,
use_sigmoid=False, use_lsgan=True, upsampling="nearest",
continue_train=False, checkpoints_dir="checkpoints/"
):
# Define input data that will be consumed by networks
self.input_A = torch.FloatTensor(
batch_size, 3, image_size, image_size
)
self.input_map = torch.FloatTensor(
batch_size, map_nc, image_size, image_size
)
norm_layer = nn.BatchNorm2d \
if norm_layer == "batch" else nn.InstanceNorm2d
# Define netD and netG
self.netG = networks.UnetGenerator(
input_nc=input_nc, output_nc=map_nc,
num_downs=num_downs, ngf=ngf,
use_dropout=use_dropout, gpu_ids=gpu_ids, norm_layer=norm_layer,
upsampling_layer=upsampling
)
self.netD = networks.NLayerDiscriminator(
input_nc=input_nc + map_nc, ndf=ndf,
n_layers=n_layers, use_sigmoid=use_sigmoid, gpu_ids=gpu_ids
)
# Transfer data to GPU
if len(gpu_ids) > 0:
self.input_A = self.input_A.cuda()
self.input_map = self.input_map.cuda()
self.netD.cuda()
self.netG.cuda()
# Initialize parameters of netD and netG
self.netG.apply(networks.weights_init)
self.netD.apply(networks.weights_init)
# Load trained netD and netG
if phase == "test" or continue_train:
netG_checkpoint_file = os.path.join(
checkpoints_dir, name, "netG_{}.pth".format(which_epoch)
)
self.netG.load_state_dict(
torch.load(netG_checkpoint_file)
)
print("Restoring netG from {}".format(netG_checkpoint_file))
if continue_train:
netD_checkpoint_file = os.path.join(
checkpoints_dir, name, "netD_{}.pth".format(which_epoch)
)
self.netD.load_state_dict(
torch.load(netD_checkpoint_file)
)
print("Restoring netD from {}".format(netD_checkpoint_file))
self.name = name
self.gpu_ids = gpu_ids
self.checkpoints_dir = checkpoints_dir
# Criterions
if phase == "train":
self.count = 0
self.lr = lr
self.lambda_D = lambda_D
self.lambda_MSE = lambda_MSE
self.image_pool = ImagePool(pool_size)
self.criterionGAN = networks.GANLoss(use_lsgan=use_lsgan)
self.criterionL1 = torch.nn.L1Loss()
self.criterionMSE = torch.nn.MSELoss() # Landmark loss
self.optimizer_G = torch.optim.Adam(
self.netG.parameters(), lr=self.lr, betas=(beta1, 0.999)
)
self.optimizer_D = torch.optim.Adam(
self.netD.parameters(), lr=self.lr, betas=(beta1, 0.999)
)
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
networks.print_network(self.netD)
print('-----------------------------------------------')
