def Trainer(opt):
# ----------------------------------------
# Initialize training parameters
# ----------------------------------------
# cudnn benchmark accelerates the network
if opt.cudnn_benchmark == True:
cudnn.benchmark = True
else:
cudnn.benchmark = False
# define network
if opt.pre_train == True:
net = network_infrared.GeneratorUNet()
utils.weights_init(net)
else:
modelname = opt.modelname
print('Load the model:', modelname)
net = torch.load(modelname)
# To device
if opt.multi_gpu == True:
net = nn.DataParallel(net)
net = net.cuda()
else:
net = net.cuda()
# Loss functions
criterion = torch.nn.L1Loss().cuda()
# Optimizers
optimizer = torch.optim.Adam(net.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
# Learning rate decrease
def adjust_learning_rate(optimizer, epoch, opt):
"""Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs"""
if opt.pre_train == False:
addition_epoch = int(opt.load_epoch)
else:
addition_epoch = 0
if opt.multi_gpu == True:
lr = opt.lr * (opt.lr_decrease_factor**(
(epoch + addition_epoch) // opt.lr_decrease_epoch))
for param_group in optimizer.module.param_groups:
param_group['lr'] = lr
else:
lr = opt.lr * (opt.lr_decrease_factor**(
(epoch + addition_epoch) // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(net, epoch, opt):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.pre_train == False:
addition_epoch = int(opt.load_epoch)
else:
addition_epoch = 0
if opt.multi_gpu == True:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net.module, 'Pre_Infrared_epoch%d_batchsize%d.pth' %
((epoch + addition_epoch), opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
epoch)
else:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net, 'Pre_Infrared_epoch%d_batchsize%d_2gammas.pth' %
((epoch + addition_epoch), opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
epoch)
# ----------------------------------------
# Initialize training dataset
# ----------------------------------------
# Network dataset txt files
# imglist = utils.text_readlines("KAIST_train.txt")
lwirlist = utils.text_readlines("KAIST_train_lwir.txt")
visiblelist = utils.text_readlines("KAIST_train_visible.txt")
# Define the dataset
trainset = dataset_infrared.InfraredDataset(opt.baseroot, lwirlist,
visiblelist)
print('The overall number of images equals to %d' % len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training and Testing
# ----------------------------------------
# Initialize start time
prev_time = time.time()
# Training loop
for epoch in range(opt.epochs):
for batch_idx, (data, target) in enumerate(dataloader):
# Load data and put it to cuda
data = data.cuda()
target = target.cuda()
# Train one iteration
optimizer.zero_grad()
output = net(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# Determine approximate time left
batches_done = epoch * len(dataloader) + batch_idx
batches_left = opt.epochs * len(dataloader) - batches_done
time_left = datetime.timedelta(seconds=batches_left *
(time.time() - prev_time))
prev_time = time.time()
# Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [L1 Loss: %.5f] time_left: %s" %
((epoch + 1), opt.epochs, batch_idx, len(dataloader),
loss.item(), time_left))
# Learning rate decrease
adjust_learning_rate(optimizer, (epoch + 1), opt)
# Save the model
save_model(net, (epoch + 1), opt)
def Trainer(opt):
# ----------------------------------------
# Initialize training parameters
# ----------------------------------------
# cudnn benchmark accelerates the network
if opt.cudnn_benchmark == True:
cudnn.benchmark = True
else:
cudnn.benchmark = False
# ResNet network
if opt.pre_train == True:
net = network_RGB_ResNet.resnet18()
utils.weights_init(net)
net = utils.load_dict(net)
else:
modelname = opt.load_type + '_epoch' + opt.load_epoch + '_batchsize' + opt.load_batchsize + '.pth'
print('Load the model:', modelname)
net = torch.load(modelname)
# To device
if opt.multi_gpu == True:
net = nn.DataParallel(net)
net = net.cuda()
else:
net = net.cuda()
# Loss functions
criterion = torch.nn.CrossEntropyLoss().cuda()
# Optimizers
optimizer = torch.optim.SGD(net.parameters(),
lr=opt.lr,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(optimizer, epoch, opt):
"""Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs"""
if opt.pre_train == False:
addition_epoch = int(opt.load_epoch)
else:
addition_epoch = 0
if opt.multi_gpu == True:
lr = opt.lr * (opt.lr_decrease_factor**(
(epoch + addition_epoch) // opt.lr_decrease_epoch))
for param_group in optimizer.module.param_groups:
param_group['lr'] = lr
else:
lr = opt.lr * (opt.lr_decrease_factor**(
(epoch + addition_epoch) // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(net, epoch, opt):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.pre_train == False:
addition_epoch = int(opt.load_epoch)
else:
addition_epoch = 0
if opt.multi_gpu == True:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net.module, '%s_IN_epoch%d_batchsize%d.pth' %
(opt.load_type, (epoch + addition_epoch), opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
epoch)
else:
if epoch % opt.checkpoint_interval == 0:
torch.save(
net, '%s_IN_epoch%d_batchsize%d.pth' %
(opt.load_type, (epoch + addition_epoch), opt.batch_size))
print('The trained model is successfully saved at epoch %d' %
epoch)
# ----------------------------------------
# Initialize training dataset
# ----------------------------------------
# Network dataset txt files
imglist = utils.text_readlines("ILSVRC2012_train_name.txt")
stringlist = utils.text_readlines("mapping_string.txt")
scalarlist = utils.text_readlines("mapping_scalar.txt")
# Define the dataset
trainset = dataset_RGB_inpainting_ImageNet.InpaintDataset(
opt, imglist, stringlist, scalarlist)
print('The overall number of images equals to %d' % len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training and Testing
# ----------------------------------------
# Initialize start time
prev_time = time.time()
# Training loop
for epoch in range(opt.epochs):
for batch_idx, (data, target) in enumerate(dataloader):
# Load data and put it to cuda
data = data.cuda()
target = target.cuda()
# Train one iteration
optimizer.zero_grad()
output = net(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# Determine approximate time left
batches_done = epoch * len(dataloader) + batch_idx
batches_left = opt.epochs * len(dataloader) - batches_done
time_left = datetime.timedelta(seconds=batches_left *
(time.time() - prev_time))
prev_time = time.time()
# Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [Cross-Entropy Loss: %.5f] time_left: %s"
% ((epoch + 1), opt.epochs, batch_idx, len(dataloader),
loss.item(), time_left))
# Learning rate decrease
adjust_learning_rate(optimizer, (epoch + 1), opt)
# Save the model
save_model(net, (epoch + 1), opt)
def __init__(self, opt):
self.opt = opt
self.imglist = utils.text_readlines(opt.namelist)
self.targetlist = utils.text_readlines(opt.targetlist)
def __init__(self, opt):
self.opt = opt
self.imglist = utils.get_jpgs(opt.baseroot)
self.stringlist = utils.text_readlines(opt.stringlist)
self.scalarlist = utils.text_readlines(opt.scalarlist)
def Train_No_GAN(opt): # w / o GAN
# ----------------------------------------
# Network training parameters
# ----------------------------------------
# cudnn benchmark
cudnn.benchmark = opt.cudnn_benchmark
# Loss functions
criterion_L1 = torch.nn.L1Loss().cuda()
# Initialize Generator
generatorNet = utils.create_generator(opt)
flownet = utils.create_pwcnet(opt)
# To device
if opt.multi_gpu:
generatorNet = nn.DataParallel(generatorNet)
generatorNet = generatorNet.cuda()
flownet = nn.DataParallel(flownet)
flownet = flownet.cuda()
else:
generatorNet = generatorNet.cuda()
flownet = flownet.cuda()
# Optimizers
optimizer_G = torch.optim.Adam(generatorNet.parameters(), lr = opt.lr_g, betas = (opt.b1, opt.b2), weight_decay = opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(opt, epoch, iteration, optimizer):
#Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs
if opt.lr_decrease_mode == 'epoch':
lr = opt.lr_g * (opt.lr_decrease_factor ** (epoch // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if opt.lr_decrease_mode == 'iter':
lr = opt.lr_g * (opt.lr_decrease_factor ** (iteration // opt.lr_decrease_iter))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(opt, epoch, iteration, len_dataset, generator):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.multi_gpu == True:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(generator.module, 'Pre_%s_epoch%d_bs%d.pth' % (opt.task, epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' % (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(generator.module, 'Pre_%s_iter%d_bs%d.pth' % (opt.task, iteration, opt.batch_size))
print('The trained model is successfully saved at iteration %d' % (iteration))
else:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(generator, 'Pre_%s_epoch%d_bs%d.pth' % (opt.task, epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' % (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(generator, 'Pre_%s_iter%d_bs%d.pth' % (opt.task, iteration, opt.batch_size))
print('The trained model is successfully saved at iteration %d' % (iteration))
# ----------------------------------------
# Network dataset
# ----------------------------------------
# Define the class list
imglist = utils.text_readlines('videocolor_linux.txt')
classlist = utils.get_dirs(opt.baseroot)
'''
imgnumber = len(imglist) - (len(imglist) % opt.batch_size)
imglist = imglist[:imgnumber]
'''
# Define the dataset
trainset = dataset.MultiFramesDataset(opt, imglist, classlist)
print('The overall number of classes:', len(trainset))
# Define the dataloader
dataloader = utils.create_dataloader(trainset, opt)
# ----------------------------------------
# Training
# ----------------------------------------
# Count start time
prev_time = time.time()
# For loop training
for epoch in range(opt.epochs):
for iteration, (in_part, out_part) in enumerate(dataloader):
# Train Generator
optimizer_G.zero_grad()
lstm_state = None
loss_flow = 0
loss_flow_long = 0
loss_L1 = 0
x_0 = in_part[0].cuda()
p_t_0 = in_part[0].cuda()
for iter_frame in range(opt.iter_frames):
# Read data
x_t = in_part[iter_frame].cuda()
y_t = out_part[iter_frame].cuda()
# Initialize the second input and compute flow loss
if iter_frame == 0:
p_t_last = torch.zeros(opt.batch_size, opt.out_channels, opt.resize_h, opt.resize_w).cuda()
elif iter_frame == 1:
x_t_last = in_part[iter_frame - 1].cuda()
p_t_last = p_t.detach()
p_t_0 = p_t.detach()
p_t_last.requires_grad = False
p_t_0.requires_grad = False
# o_t_last_2_t range is [-20, +20]
o_t_last_2_t = pwcnet.PWCEstimate(flownet, x_t, x_t_last)
x_t_warp = pwcnet.PWCNetBackward((x_t_last + 1) / 2, o_t_last_2_t)
# y_t_warp range is [0, 1]
p_t_warp = pwcnet.PWCNetBackward((p_t_last + 1) / 2, o_t_last_2_t)
else:
x_t_last = in_part[iter_frame - 1].cuda()
p_t_last = p_t.detach()
p_t_last.requires_grad = False
# o_t_last_2_t o_t_first_2_t range is [-20, +20]
o_t_last_2_t = pwcnet.PWCEstimate(flownet, x_t, x_t_last)
o_t_first_2_t = pwcnet.PWCEstimate(flownet,x_t, x_0)
# y_t_warp, y_t_warp_long range is [0, 1]
x_t_warp = pwcnet.PWCNetBackward((x_t_last + 1) / 2, o_t_last_2_t)
p_t_warp = pwcnet.PWCNetBackward((p_t_last + 1) / 2, o_t_last_2_t)
x_t_warp_long = pwcnet.PWCNetBackward((x_0 + 1) / 2, o_t_first_2_t)
p_t_warp_long = pwcnet.PWCNetBackward((p_t_0 + 1) / 2, o_t_first_2_t)
# Generator output
p_t, lstm_state = generatorNet(x_t, p_t_last, lstm_state)
lstm_state = utils.repackage_hidden(lstm_state)
if iter_frame == 1:
mask_flow = torch.exp( -opt.mask_para * torch.sum((x_t + 1) / 2 - x_t_warp, dim=1).pow(2) ).unsqueeze(1)
loss_flow += criterion_L1(mask_flow * (p_t + 1) / 2, mask_flow * p_t_warp)
elif iter_frame > 1:
mask_flow = torch.exp( -opt.mask_para * torch.sum((x_t + 1) / 2 - x_t_warp, dim=1).pow(2) ).unsqueeze(1)
loss_flow += criterion_L1(mask_flow * (p_t + 1) / 2, mask_flow * p_t_warp)
mask_flow_long = torch.exp( -opt.mask_para * torch.sum((x_t + 1) / 2 - x_t_warp_long, dim=1).pow(2) ).unsqueeze(1)
loss_flow_long += criterion_L1(mask_flow_long * (p_t + 1) / 2, mask_flow_long * p_t_warp_long)
# Pixel-level loss
loss_L1 += criterion_L1(p_t, y_t)
# Overall Loss and optimize
loss = loss_L1 + opt.lambda_flow * loss_flow + opt.lambda_flow_long * loss_flow_long
loss.backward()
optimizer_G.step()
# Determine approximate time left
iters_done = epoch * len(dataloader) + iteration
iters_left = opt.epochs * len(dataloader) - iters_done
time_left = datetime.timedelta(seconds = iters_left * (time.time() - prev_time))
prev_time = time.time()
# Print log
print("\r[Epoch %d/%d] [Batch %d/%d] [L1 Loss: %.4f] [Flow Loss Short: %.8f] [Flow Loss Long: %.8f] Time_left: %s" %
((epoch + 1), opt.epochs, iteration, len(dataloader), loss_L1.item(), loss_flow.item(), loss_flow_long.item(), time_left))
# Save model at certain epochs or iterations
save_model(opt, (epoch + 1), (iters_done + 1), len(dataloader), generatorNet)
# Learning rate decrease at certain epochs
adjust_learning_rate(opt, (epoch + 1), (iters_done + 1), optimizer_G)
def Pre_train_single(opt):
# ----------------------------------------
# Network training parameters
# ----------------------------------------
print("Pre_train_single")
# cudnn benchmark
cudnn.benchmark = opt.cudnn_benchmark
# Loss functions
criterion_L1 = torch.nn.L1Loss().cuda()
criterion_MSE = torch.nn.MSELoss().cuda()
# Initialize Generator
generatorNet = utils.create_generator(opt)
# To device
if opt.multi_gpu:
generatorNet = nn.DataParallel(generatorNet)
generatorNet = generatorNet.cuda()
else:
generatorNet = generatorNet.cuda()
# Optimizers
optimizer_G = torch.optim.Adam(generatorNet.parameters(), lr = opt.lr_g, betas = (opt.b1, opt.b2), weight_decay = opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(opt, epoch, iteration, optimizer):
#Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs
if opt.lr_decrease_mode == 'epoch':
lr = opt.lr_g * (opt.lr_decrease_factor ** (epoch // opt.lr_decrease_epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if opt.lr_decrease_mode == 'iter':
lr = opt.lr_g * (opt.lr_decrease_factor ** (iteration // opt.lr_decrease_iter))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Save the model if pre_train == True
def save_model(opt, epoch, iteration, len_dataset, generator):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.multi_gpu == True:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(generator.module, 'Pre_%s_epoch%d_bs%d.pth' % (opt.task, epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' % (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(generator.module, 'Pre_%s_iter%d_bs%d.pth' % (opt.task, iteration, opt.batch_size))
print('The trained model is successfully saved at iteration %d' % (iteration))
else:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch == 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(generator, 'Pre_%s_epoch%d_bs%d.pth' % (opt.task, epoch, opt.batch_size))
print('The trained model is successfully saved at epoch %d' % (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(generator, 'Pre_%s_iter%d_bs%d.pth' % (opt.task, iteration, opt.batch_size))
print('The trained model is successfully saved at iteration %d' % (iteration))
# ----------------------------------------
# Network dataset
# ----------------------------------------
# Define the class list
imglist = utils.text_readlines('ILSVRC2012_train_sal_name.txt')[:1272480]
# Define the dataset
trainset = dataset.ColorizationDataset(opt, imglist)
print('The overall number of classes:', len(trainset))
# Define the dataloader
dataloader = DataLoader(trainset, batch_size = opt.batch_size, shuffle = True, num_workers = opt.num_workers, pin_memory = True)
# ----------------------------------------
# Training
# ----------------------------------------
# Tensor type
Tensor = torch.cuda.FloatTensor
# Count start time
prev_time = time.time()
# For loop training
# For loop training
for epoch in range(opt.epochs):
for iteration, (x_t, y_t) in enumerate(dataloader):
# Train Generator
optimizer_G.zero_grad()
lstm_state = None
x_t = x_t.cuda()
y_t = y_t.cuda()
valid = Tensor(np.ones((x_t.shape[0], 1, 30, 30)))
p_t_last = torch.zeros(opt.batch_size, opt.out_channels, opt.resize_h, opt.resize_w).cuda()
# Generator output
p_t, lstm_state = generatorNet(x_t, p_t_last, lstm_state)
# Pixel-level loss
loss_L1 = criterion_L1(p_t, y_t)
# Overall Loss and optimize
loss = loss_L1
loss.backward()
optimizer_G.step()
# Determine approximate time left
iters_done = epoch * len(dataloader) + iteration
iters_left = opt.epochs * len(dataloader) - iters_done
time_left = datetime.timedelta(seconds = iters_left * (time.time() - prev_time))
prev_time = time.time()
# Print log
print("\r[Epoch %d/%d] [Batch %d/%d] [L1 Loss: %.4f] Time_left: %s" %
((epoch + 1), opt.epochs, iteration, len(dataloader), loss_L1.item(), time_left))
# Save model at certain epochs or iterations
save_model(opt, (epoch + 1), (iters_done + 1), len(dataloader), generatorNet)
# Learning rate decrease at certain epochs
adjust_learning_rate(opt, (epoch + 1), (iters_done + 1), optimizer_G)