def main():
opt = TestOptions().parse(save=False)
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
if not opt.engine and not opt.onnx:
model = create_model(opt)
if opt.data_type == 16:
model.half()
elif opt.data_type == 8:
model.type(torch.uint8)
if opt.verbose:
print(model)
else:
from run_engine import run_trt_engine, run_onnx
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
if opt.data_type == 16:
data['label'] = data['label'].half()
data['inst'] = data['inst'].half()
elif opt.data_type == 8:
data['label'] = data['label'].uint8()
data['inst'] = data['inst'].uint8()
if opt.export_onnx:
print("Exporting to ONNX: ", opt.export_onnx)
assert opt.export_onnx.endswith(
"onnx"), "Export model file should end with .onnx"
torch.onnx.export(model, [data['label'], data['inst']],
opt.export_onnx,
verbose=True)
exit(0)
minibatch = 1
if opt.engine:
generated = run_trt_engine(opt.engine, minibatch,
[data['label'], data['inst']])
elif opt.onnx:
generated = run_onnx(opt.onnx, opt.data_type, minibatch,
[data['label'], data['inst']])
else:
generated = model.inference(data['label'], data['inst'],
data['image'])
visuals = OrderedDict([
('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
('synthesized_image', util.tensor2im(generated.data[0]))
])
img_path = data['path']
print('process image... %s' % img_path)
visualizer.save_images(webpage, visuals, img_path)
webpage.save()
num_batch = 5
# load options
opt = TestEncoderDecoderOptions_V2().parse()
opt.batch_size = batch_size
train_opt = io.load_json(os.path.join('checkpoints', opt.id, 'train_opt.json'))
preserved_opt = {'gpu_ids', 'batch_size', 'is_train'}
for k, v in train_opt.iteritems():
if k in opt and (k not in preserved_opt):
setattr(opt, k, v)
# create model
model = EncoderDecoderFramework_DFN()
model.initialize(opt)
# create data loader
val_loader_iter = iter(CreateDataLoader(opt, split='test'))
# create viusalizer
visualizer = GANVisualizer_V3(opt)
for i in range(num_batch):
print('[%s] test edge transfer: %d / %d' % (opt.id, i+1, num_batch))
data = val_loader_iter.next()
imgs_title = data['img'].expand(batch_size, 3, opt.fine_size, opt.fine_size).cpu()
for name in ['img', 'seg_map', 'seg_mask', 'pose_map', 'edge_map', 'color_map']:
# print('%s: %s' % (name, data[name].size()))
new_1 = torch.cat([data[name]] * batch_size, dim=0).contiguous()
new_2 = new_1.view(batch_size, batch_size, -1).transpose(0, 1).contiguous().view(new_1.size())
new_3 = torch.cat([data[name+'_def']] * batch_size, dim=0).contiguous()
new_3 = new_3.view(batch_size, batch_size, -1).transpose(0, 1).contiguous().view(new_1.size())
data[name] = new_2
def main():
opt.input_image_root = './data/%s_%s_%s' % (opt.dataset_name, opt.model_name, opt.which_epoch)
opt.real_image_root = '../data/%s/train/face_img' % opt.dataset_name
opt.label_root = '../data/%s/train/face_label' % opt.dataset_name
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
start_epoch, epoch_iter = 1, 0
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
model = create_model(opt)
# model = model.cuda()
visualizer = Visualizer(opt)
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = epoch_iter % dataset_size
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == display_delta
############## Forward Pass ######################
losses, generated = model(Variable(data['label']), Variable(data['inst']),
Variable(data['image']), Variable(data['feat']), infer=save_fake)
# sum per device losses
losses = [torch.mean(x) if not isinstance(x, int) else x for x in losses]
loss_dict = dict(zip(model.loss_names, losses))
# calculate final loss scalar
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0)
############### Backward Pass ####################
# update generator weights
model.optimizer_G.zero_grad()
loss_G.backward()
model.optimizer_G.step()
# update discriminator weights
model.optimizer_D.zero_grad()
loss_D.backward()
model.optimizer_D.step()
############## Display results and errors ##########
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {k: v.data if not isinstance(v, int) else v for k, v in loss_dict.items()}
t = (time.time() - iter_start_time) / opt.batchSize
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
### display output images
if save_fake:
input_image = data['label'].detach()[0][10:,:,:]
input_label = data['label'].detach()[0][:10,:,:]
generated_image = generated.detach()[0]
visuals = OrderedDict([('input_label', util.tensor2label(input_label, opt.label_nc)),
('input_image', util.tensor2im(input_image)),
('synthesized_image', util.tensor2im(generated_image)),
('real_image', util.tensor2im(data['image'][0]))])
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
if total_steps % opt.save_latest_freq == save_delta:
print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps))
model.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d')
if epoch_iter >= dataset_size:
break
# end of epoch
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
### save model for this epoch
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps))
model.save('latest')
model.save(epoch)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
### instead of only training the local enhancer, train the entire network after certain iterations
if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
model.update_fixed_params()
### linearly decay learning rate after certain iterations
if epoch > opt.niter:
model.update_learning_rate()
torch.cuda.empty_cache()
def create_descriptor_vunet():
'''
use variational unet feature as descriptor
'''
import torch
from models.vunet_pose_transfer_model import VUnetPoseTransferModel
from data.data_loader import CreateDataLoader
from options.pose_transfer_options import TestPoseTransferOptions
# load image info
image_info = io.load_json('temp/patch_matching/label/image_info.json')
# load model
# opt = TestPoseTransferOptions().parse(ord_str = '--which_model_T vunet --id 7.5 --gpu_ids 0,1,2,3 --batch_size 16', save_to_file = False, display = False, set_gpu = True)
opt = TestPoseTransferOptions().parse(
ord_str=
'--which_model_T vunet --id 7.9 --gpu_ids 0,1,2,3 --batch_size 16',
save_to_file=False,
display=False,
set_gpu=True)
train_opt = io.load_json(
os.path.join('checkpoints', opt.id, 'train_opt.json'))
for k, v in train_opt.iteritems():
if k in opt and (k not in {'gpu_ids', 'is_train'}):
setattr(opt, k, v)
model = VUnetPoseTransferModel()
model.initialize(opt)
# create data set
val_loader = CreateDataLoader(opt, split='test')
id_list = [[sid_1, sid_2]
for sid_1, sid_2 in zip(image_info['id_1'], image_info['id_2'])]
val_loader.dataset.id_list = id_list
# extract descriptor
print('extracing descriptors ...')
desc = {
'img_1': [],
'img_2': [],
'feat1_1': [],
'feat1_2': [],
'feat2_1': [],
'feat2_2': [],
}
for i, data in enumerate(tqdm.tqdm(val_loader)):
model.set_input(data)
appr_1 = model.get_appearance(model.opt.appearance_type, index='1')
appr_2 = model.get_appearance(model.opt.appearance_type, index='2')
pose_1 = model.get_pose(model.opt.pose_type, index='1')
pose_2 = model.get_pose(model.opt.pose_type, index='2')
# desc_1 (ref)
output, ps, qs, ds = model.netT(appr_1,
pose_1,
pose_1,
mode='transfer',
output_feature=True)
output = model.parse_output(output, model.opt.output_type)
img = torch.nn.functional.tanh(output['image'])
desc['img_1'].append(img.detach().cpu())
desc['feat1_1'].append(ds[-1].detach().cpu())
desc['feat2_1'].append(ds[-2].detach().cpu())
# desc_2 (tar)
output, ps, qs, ds = model.netT(appr_1,
pose_1,
pose_2,
mode='transfer',
output_feature=True)
output = model.parse_output(output, model.opt.output_type)
img = torch.nn.functional.tanh(output['image'])
desc['img_2'].append(img.detach().cpu())
desc['feat1_2'].append(ds[-1].detach().cpu())
desc['feat2_2'].append(ds[-2].detach().cpu())
for k, v in desc.iteritems():
desc[k] = torch.cat(v).numpy().transpose(0, 2, 3, 1)
# save descriptor
data_dict_img = {
'desc_1': desc['img_1'],
'desc_2': desc['img_2'],
'name': 'gen_vunet_img'
}
data_dict_feat1 = {
'desc_1': desc['feat1_1'],
'desc_2': desc['feat1_2'],
'name': 'gen_vunet_feat1'
}
data_dict_feat2 = {
'desc_1': desc['feat2_1'],
'desc_2': desc['feat2_2'],
'name': 'gen_vunet_feat2'
}
scipy.io.matlab.savemat(
'temp/patch_matching/descriptor/desc_gen_vunet-7.9_img.mat',
data_dict_img)
scipy.io.matlab.savemat(
'temp/patch_matching/descriptor/desc_gen_vunet-7.9_feat1.mat',
data_dict_feat1)
scipy.io.matlab.savemat(
'temp/patch_matching/descriptor/desc_gen_vunet-7.9_feat2.mat',
data_dict_feat2)
# visualize desc_img
vis_dir = 'temp/patch_matching/descriptor/vis_gen_vunet_img/'
io.mkdir_if_missing(vis_dir)
imgs_1 = (desc['img_1'] * 127.5 + 127.5).astype(np.uint8)
imgs_2 = (desc['img_2'] * 127.5 + 127.5).astype(np.uint8)
for i in range(imgs_1.shape[0]):
imageio.imwrite(vis_dir + '%d_1.jpg' % i, imgs_1[i])
imageio.imwrite(vis_dir + '%d_2.jpg' % i, imgs_2[i])
import sys
from options.train_options import TrainOptions
opt = TrainOptions().parse() # set CUDA_VISIBLE_DEVICES before import torch
from data.data_loader import CreateDataLoader
from models.models import create_model
dataset_root = "/phoenix/S6/zl548/"
test_list_dir_l = '/phoenix/S6/zl548/MegaDpeth_code/test_list/landscape/'
input_height = 240
input_width = 320
is_flipped = False
shuffle = False
test_data_loader_l = CreateDataLoader(dataset_root, test_list_dir_l,
input_height, input_width, is_flipped,
shuffle)
test_dataset_l = test_data_loader_l.load_data()
test_dataset_size_l = len(test_data_loader_l)
print('========================= test images = %d' % test_dataset_size_l)
test_list_dir_p = '/phoenix/S6/zl548/MegaDpeth_code/test_list/portrait/'
input_height = 320
input_width = 240
test_data_loader_p = CreateDataLoader(dataset_root, test_list_dir_p,
input_height, input_width, is_flipped,
shuffle)
test_dataset_p = test_data_loader_p.load_data()
test_dataset_size_p = len(test_data_loader_p)
print('========================= test images = %d' % test_dataset_size_p)
model = create_model(opt)
import time
from options.train_options import TrainOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
import copy
import numpy as np
import os
opt = TrainOptions().parse()
data_loader = CreateDataLoader(opt)
opt_copy = copy.deepcopy(opt)
opt_copy.dataset_mode = 'aligned'
opt_copy.dataroot = opt.dataroot_aligned
opt_copy.resize_or_crop = 'resize_and_crop'
paired_data_loader = CreateDataLoader(opt_copy)
aligned_dataset = paired_data_loader.load_data()
dataset = data_loader.load_data()
dataset_size = len(data_loader) + len(paired_data_loader)
print('#training images = %d' % dataset_size)
#model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = 0
for epoch in range(1): #opt.epoch_count, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
epoch_iter = 0
from models.models import create_model
from util.visualizer import Visualizer
from pdb import set_trace as st
from util import html
from data.process_videos import avi2pngs
opt = TrainOptions().parse()
opt.nThreads = 1 # test code only supports nThreads = 1
opt.continue_train = True
opt.dataroot = opt.trainroot
opt.file_list = opt.train_file_list
avi2pngs(opt)
opt.dataroot = opt.dataroot + '/split'
training_data_loader = CreateDataLoader(opt)
train_dataset = training_data_loader.load_data()
train_dataset_size = len(training_data_loader)
print('#training samples = %d' % train_dataset_size)
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
opt.dataroot = opt.testroot
opt.file_list = opt.test_file_list
avi2pngs(opt)
opt.dataroot = opt.dataroot + '/split'
testing_data_loader = CreateDataLoader(opt)
test_dataset = testing_data_loader.load_data()
test_dataset_size = len(testing_data_loader)
def main():
opt = TrainOptions().parse()
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
# read pix2pix/PAN moodel
if opt.model == 'pix2pix':
assert (opt.dataset_mode == 'aligned')
from models.pix2pix_model import Pix2PixModel
model = Pix2PixModel()
model.initialize(opt)
elif opt.model == 'pan':
from models.pan_model import PanModel
model = PanModel()
model.initialize(opt)
total_steps = 0
batch_size = opt.batchSize
print_freq = opt.print_freq
epoch_count = opt.epoch_count
niter = opt.niter
niter_decay = opt.niter_decay
display_freq = opt.display_freq
save_latest_freq = opt.save_latest_freq
save_epoch_freq = opt.save_epoch_freq
for epoch in range(epoch_count, niter + niter_decay + 1):
epoch_start_time = time.time()
epoch_iter = 0
for i, data in enumerate(dataset):
# data --> (1, 3, 256, 256)
iter_start_time = time.time()
total_steps += batch_size
epoch_iter += batch_size
model.set_input(data)
model.optimize_parameters()
if total_steps % print_freq == 0:
errors = model.get_current_errors()
t = (time.time() - iter_start_time) / batch_size
message = '(epoch: %d, iters: %d, time: %.3f) ' % (
epoch, epoch_iter, t)
for k, v in errors.items():
message += '%s: %.3f ' % (k, v)
print(message)
# save latest weights
if total_steps % save_latest_freq == 0:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch, total_steps))
model.save('latest')
# save weights periodicaly
if epoch % save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
model.save('latest')
model.save(epoch)
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, niter + niter_decay, time.time() - epoch_start_time))
model.update_learning_rate()
import tqdm
import time
from collections import OrderedDict
# parse and save options
parser = TrainPoseTransferOptions()
opt = parser.parse()
parser.save()
# create model
model = PoseTransferModel()
model.initialize(opt)
# save terminal line
io.save_str_list([' '.join(sys.argv)],
os.path.join(model.save_dir, 'order_line.txt'))
# create data loader
train_loader = CreateDataLoader(opt, split='train')
val_loader = CreateDataLoader(
opt, split='test' if not opt.small_val_set else 'test_small')
# create visualizer
visualizer = Visualizer(opt)
logdir = os.path.join('logs', opt.id)
if not os.path.exists(logdir):
os.makedirs(logdir)
writer = tensorboardX.SummaryWriter(logdir)
# set "saving best"
best_info = {'meas': 'SSIM', 'type': 'max', 'best_value': 0, 'best_epoch': -1}
# set continue training
if not opt.resume_train:
import sys
import time
import numpy as np
import torch.backends.cudnn as cudnn
from tqdm import tqdm
from data.data_loader import CreateDataLoader
from models.models import create_model
from options.train_options import TrainOptions
cudnn.benchmark = True
opt = TrainOptions().parse()
ROOT = '/data1'
testDataLoader = CreateDataLoader(opt, {'mode':'Test', 'labelFn':'in5008.txt', 'rootPath':ROOT, 'subDir':'eval'})
#testDataLoader = CreateDataLoader(opt, {'mode':'Train', 'labelFn':'in5008.txt', 'rootPath':ROOT, 'subDir':'train'})
testDataset = testDataLoader.load_data()
print('#testing images = %d' % len(testDataLoader))
model = create_model(opt)
totalWer = list()
totalCer = list()
for i, data in tqdm(enumerate(testDataset), total=len(testDataset)):
model.set_input(data)
results = model.test()
totalWer.append(results['wer'])
totalCer.append(results['cer'])
def train_pose2vid(target_dir, run_name, temporal_smoothing=False):
import src.config.train_opt as opt
opt = update_opt(opt, target_dir, run_name, temporal_smoothing)
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.json')
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
if opt.load_pretrain != '':
with open(iter_path, 'r') as f:
iter_json = json.load(f)
else:
iter_json = {'start_epoch': 1, 'epoch_iter': 0}
start_epoch = iter_json['start_epoch']
epoch_iter = iter_json['epoch_iter']
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
model = create_model(opt)
model = model.to(device)
visualizer = Visualizer(opt)
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = epoch_iter % dataset_size
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == display_delta
############## Forward Pass ######################
if temporal_smoothing:
losses, generated = model(Variable(data['label']),
Variable(data['inst']),
Variable(data['image']),
Variable(data['feat']),
Variable(data['previous_label']),
Variable(data['previous_image']),
infer=save_fake)
else:
losses, generated = model(Variable(data['label']),
Variable(data['inst']),
Variable(data['image']),
Variable(data['feat']),
infer=save_fake)
# sum per device losses
losses = [
torch.mean(x) if not isinstance(x, int) else x for x in losses
]
loss_dict = dict(zip(model.loss_names, losses))
# calculate final loss scalar
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + loss_dict.get(
'G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0)
############### Backward Pass ####################
# update generator weights
model.optimizer_G.zero_grad()
loss_G.backward()
model.optimizer_G.step()
# update discriminator weights
model.optimizer_D.zero_grad()
loss_D.backward()
model.optimizer_D.step()
############## Display results and errors ##########
print(f"Epoch {epoch} batch {i}:")
print(f"loss_D: {loss_D}, loss_G: {loss_G}")
print(
f"loss_D_fake: {loss_dict['D_fake']}, loss_D_real: {loss_dict['D_real']}"
)
print(
f"loss_G_GAN {loss_dict['G_GAN']}, loss_G_GAN_Feat: {loss_dict.get('G_GAN_Feat', 0)}, loss_G_VGG: {loss_dict.get('G_VGG', 0)}\n"
)
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {
k: v.item() if not isinstance(v, int) else v
for k, v in loss_dict.items()
}
# errors = {k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items()}
t = (time.time() - iter_start_time) / opt.batchSize
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
### display output images
if save_fake:
visuals = OrderedDict([
('input_label',
util.tensor2label(data['label'][0], opt.label_nc)),
('synthesized_image', util.tensor2im(generated.data[0])),
('real_image', util.tensor2im(data['image'][0]))
])
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
if total_steps % opt.save_latest_freq == save_delta:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch, total_steps))
model.save('latest')
iter_json['start_epoch'] = epoch
iter_json['epoch_iter'] = epoch_iter
with open(iter_path, 'w') as f:
json.dump(iter_json, f)
if epoch_iter >= dataset_size:
break
# end of epoch
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
### save model for this epoch
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
model.save('latest')
model.save(epoch)
iter_json['start_epoch'] = epoch + 1
iter_json['epoch_iter'] = 0
with open(iter_path, 'w') as f:
json.dump(iter_json, f)
### instead of only training the local enhancer, train the entire network after certain iterations
if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
model.update_fixed_params()
### linearly decay learning rate after certain iterations
if epoch > opt.niter:
model.update_learning_rate()
torch.cuda.empty_cache()
def train():
opt = TrainOptions().parse()
if opt.distributed:
init_dist()
print('batch size per GPU: %d' % opt.batchSize)
torch.backends.cudnn.benchmark = True
### setup dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
pose = 'pose' in opt.dataset_mode
### setup trainer
trainer = Trainer(opt, data_loader)
### setup models
model, flowNet, [optimizer_G,
optimizer_D] = create_model(opt, trainer.start_epoch)
flow_gt = conf_gt = [None] * 2
for epoch in range(trainer.start_epoch, opt.niter + opt.niter_decay + 1):
if opt.distributed:
dataset.sampler.set_epoch(epoch)
trainer.start_of_epoch(epoch, model, data_loader)
n_frames_total, n_frames_load = data_loader.dataset.n_frames_total, opt.n_frames_per_gpu
for idx, data in enumerate(dataset, start=trainer.epoch_iter):
data = trainer.start_of_iter(data)
if not opt.no_flow_gt:
data_list = [
data['tgt_label'], data['ref_label']
] if pose else [data['tgt_image'], data['ref_image']]
flow_gt, conf_gt = flowNet(data_list, epoch)
data_list = [
data['tgt_label'], data['tgt_image'], flow_gt, conf_gt
]
data_ref_list = [data['ref_label'], data['ref_image']]
data_prev = [None, None, None]
############## Forward Pass ######################
for t in range(0, n_frames_total, n_frames_load):
data_list_t = get_data_t(data_list, n_frames_load,
t) + data_ref_list + data_prev
d_losses = model(data_list_t, mode='discriminator')
d_losses = loss_backward(opt, d_losses, optimizer_D, 1)
g_losses, generated, data_prev = model(
data_list_t, save_images=trainer.save, mode='generator')
g_losses = loss_backward(opt, g_losses, optimizer_G, 0)
loss_dict = dict(
zip(model.module.lossCollector.loss_names,
g_losses + d_losses))
if trainer.end_of_iter(loss_dict,
generated + data_list + data_ref_list,
model):
break
trainer.end_of_epoch(model)
def main():
with open('./train/train_opt.pkl', mode='rb') as f:
opt = pickle.load(f)
opt.checkpoints_dir = './checkpoints/'
opt.dataroot = './train'
opt.no_flip = True
opt.label_nc = 0
opt.batchSize = 2
print(opt)
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
start_epoch, epoch_iter = 1, 0
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
best_loss = 999999
epoch_loss = 9999999999
model = create_model(opt)
model = model.cuda()
visualizer = Visualizer(opt)
#niter = 20,niter_decay = 20
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = epoch_iter % dataset_size
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == display_delta
############## Forward Pass ######################
losses, generated = model(Variable(data['label']), Variable(data['inst']),
Variable(data['image']), Variable(data['feat']), infer=save_fake)
# sum per device losses
losses = [torch.mean(x) if not isinstance(x, int) else x for x in losses]
loss_dict = dict(zip(model.loss_names, losses))
# calculate final loss scalar
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + loss_dict.get('G_GAN_Feat', 0) + loss_dict.get('G_VGG', 0)
loss_DG = loss_D + loss_G
############### Backward Pass ####################
# update generator weights
model.optimizer_G.zero_grad()
loss_G.backward()
model.optimizer_G.step()
# update discriminator weights
model.optimizer_D.zero_grad()
loss_D.backward()
model.optimizer_D.step()
# call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"])
############## Display results and errors ##########
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {k: v.data[0] if not isinstance(v, int) else v for k, v in loss_dict.items()}
t = (time.time() - iter_start_time) / opt.batchSize
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
### display output images
if save_fake:
visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
('synthesized_image', util.tensor2im(generated.data[0])),
('real_image', util.tensor2im(data['image'][0]))])
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
if total_steps % opt.save_latest_freq == save_delta and loss_DG<best_loss:
best_loss = loss_DG
print('saving the latest model (epoch %d, total_steps %d ,total loss %g)' % (epoch, total_steps,loss_DG.item()))
model.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d')
if epoch_iter >= dataset_size:
break
# end of epoch
iter_end_time = time.time()
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
### save model for this epoch
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d ' % (epoch, total_steps))
model.save('latest')
model.save(epoch)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
### instead of only training the local enhancer, train the entire network after certain iterations
if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
model.update_fixed_params()
### linearly decay learning rate after certain iterations
if epoch > opt.niter:
model.update_learning_rate()
torch.cuda.empty_cache()
import time
from options.train_options import TrainOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
import pdb
pdb.set_trace()
opt = TrainOptions().parse()
# data for training
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
# data for test
from copy import deepcopy
test_opt = deepcopy(opt)
test_opt.phase = 'test'
test_opt.nThreads = 1 # test code only supports nThreads = 1
test_opt.batchSize = 1 # test code only supports batchSize = 1
test_opt.serial_batches = True # no shuffle
test_opt.no_flip = True # no flip
test_data_loader = CreateDataLoader(test_opt)
test_dataset = test_data_loader.load_data()
previous_score = 0.0
#########
model = create_model(opt)
visualizer = Visualizer(opt)
from tensorboardX import SummaryWriter
from util.preprocess import preprocess_for_train
from util.extract_key_bases import extract_key_bases
if __name__ == "__main__":
opt = TrainOptions().parse()
# 对数据进行预处理
print('Pre-processing datasets for train ...')
start = time.time()
preprocess_for_train(opt.dataset_root)
end = time.time()
print('Pre-process for train completed !')
print('Pre-process time : %d min' % int((end - start) / 60))
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
#create validation set data loader if validation_on option is set
if opt.validation_on:
#temperally set to val to load val data
opt.mode = 'val'
data_loader_val = CreateDataLoader(opt)
dataset_val = data_loader_val.load_data()
dataset_size_val = len(data_loader_val)
print('#validation images = %d' % dataset_size_val)
opt.mode = 'train' #set it back
writer = SummaryWriter(comment=opt.name)
import time
from options.train_options import TrainOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
import cv2
import copy
import pdb
opt = TrainOptions().parse()
opt2 = copy.copy(opt)
# change some opts
opt2.dataroot = opt.dataroot_more
#opt.which_direction = 'BtoA' # ensure the common domain A is photo
#opt2.which_direction = 'BtoA'
data_loader = CreateDataLoader(opt)
data_loader2 = CreateDataLoader(opt2)
dataset = data_loader.load_data()
dataset2 = data_loader2.load_data()
dataset_size = min(len(data_loader), len(data_loader2))
#dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = 0
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
epoch_iter = 0
def train():
opt = TrainOptions().parse()
if opt.debug:
opt.display_freq = 1
opt.print_freq = 1
opt.nThreads = 1
### initialize dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training videos = %d' % dataset_size)
### initialize models
models = create_model(opt)
modelG, modelD, flowNet, optimizer_G, optimizer_D, optimizer_D_T = create_optimizer(
opt, models)
### set parameters
n_gpus, tG, tD, tDB, s_scales, t_scales, input_nc, output_nc, \
start_epoch, epoch_iter, print_freq, total_steps, iter_path = init_params(opt, modelG, modelD, dataset_size)
visualizer = Visualizer(opt)
### real training starts here
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
for idx, data in enumerate(dataset, start=epoch_iter):
if total_steps % print_freq == 0:
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == 0
n_frames_total, n_frames_load, t_len = data_loader.dataset.init_data_params(
data, n_gpus, tG)
fake_B_last, frames_all = data_loader.dataset.init_data(t_scales)
for i in range(0, n_frames_total, n_frames_load):
input_A, input_B, inst_A = data_loader.dataset.prepare_data(
data, i, input_nc, output_nc)
###################################### Forward Pass ##########################
####### generator
fake_B, fake_B_raw, flow, weight, real_A, real_Bp, fake_B_last = modelG(
input_A, input_B, inst_A, fake_B_last)
####### discriminator
### individual frame discriminator
real_B_prev, real_B = real_Bp[:, :
-1], real_Bp[:,
1:] # the collection of previous and current real frames
flow_ref, conf_ref = flowNet(
real_B, real_B_prev) # reference flows and confidences
#flow_ref, conf_ref = util.remove_dummy_from_tensor([flow_ref, conf_ref])
fake_B_prev = modelG.module.compute_fake_B_prev(
real_B_prev, fake_B_last, fake_B)
losses = modelD(
0,
reshape([
real_B, fake_B, fake_B_raw, real_A, real_B_prev,
fake_B_prev, flow, weight, flow_ref, conf_ref
]))
losses = [
torch.mean(x) if x is not None else 0 for x in losses
]
loss_dict = dict(zip(modelD.module.loss_names, losses))
### temporal discriminator
# get skipped frames for each temporal scale
frames_all, frames_skipped = modelD.module.get_all_skipped_frames(frames_all, \
real_B, fake_B, flow_ref, conf_ref, t_scales, tD, n_frames_load, i, flowNet)
# run discriminator for each temporal scale
loss_dict_T = []
for s in range(t_scales):
if frames_skipped[0][s] is not None:
losses = modelD(s + 1, [
frame_skipped[s]
for frame_skipped in frames_skipped
])
losses = [
torch.mean(x) if not isinstance(x, int) else x
for x in losses
]
loss_dict_T.append(
dict(zip(modelD.module.loss_names_T, losses)))
# collect losses
loss_G, loss_D, loss_D_T, t_scales_act = modelD.module.get_losses(
loss_dict, loss_dict_T, t_scales)
###################################### Backward Pass #################################
# update generator weights
loss_backward(opt, loss_G, optimizer_G)
# update individual discriminator weights
loss_backward(opt, loss_D, optimizer_D)
# update temporal discriminator weights
for s in range(t_scales_act):
loss_backward(opt, loss_D_T[s], optimizer_D_T[s])
if i == 0:
fake_B_first = fake_B[
0, 0] # the first generated image in this sequence
if opt.debug:
call([
"nvidia-smi", "--format=csv",
"--query-gpu=memory.used,memory.free"
])
############## Display results and errors ##########
### print out errors
if total_steps % print_freq == 0:
t = (time.time() - iter_start_time) / print_freq
errors = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict.items()
}
for s in range(len(loss_dict_T)):
errors.update({
k + str(s):
v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict_T[s].items()
})
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
### display output images
if save_fake:
visuals = util.save_all_tensors(opt, real_A, fake_B,
fake_B_first, fake_B_raw,
real_B, flow_ref, conf_ref,
flow, weight, modelD)
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
save_models(opt, epoch, epoch_iter, total_steps, visualizer,
iter_path, modelG, modelD)
if epoch_iter > dataset_size - opt.batchSize:
epoch_iter = 0
break
# end of epoch
iter_end_time = time.time()
visualizer.vis_print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
### save model for this epoch and update model params
save_models(opt,
epoch,
epoch_iter,
total_steps,
visualizer,
iter_path,
modelG,
modelD,
end_of_epoch=True)
update_models(opt, epoch, modelG, modelD, data_loader)
def train():
opt = TrainOptions().parse()
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
if opt.continue_train:
try:
start_epoch, epoch_iter = np.loadtxt(iter_path,
delimiter=',',
dtype=int)
except:
start_epoch, epoch_iter = 1, 0
# compute resume lr
if start_epoch > opt.niter:
lrd_unit = opt.lr / opt.niter_decay
resume_lr = opt.lr - (start_epoch - opt.niter) * lrd_unit
opt.lr = resume_lr
print('Resuming from epoch %d at iteration %d' %
(start_epoch, epoch_iter))
else:
start_epoch, epoch_iter = 1, 0
opt.print_freq = lcm(opt.print_freq, opt.batchSize)
if opt.debug:
opt.display_freq = 2
opt.print_freq = 2
opt.niter = 3
opt.niter_decay = 0
opt.max_dataset_size = 1
opt.valSize = 1
## Loading data
# train data
data_loader = CreateDataLoader(opt, isVal=False)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('# training images = %d' % dataset_size)
# validation data
data_loader = CreateDataLoader(opt, isVal=True)
valset = data_loader.load_data()
print('# validation images = %d' % len(data_loader))
## Loading model
model = create_model(opt)
visualizer = Visualizer(opt)
if opt.fp16:
from apex import amp
model, [optimizer_G, optimizer_D
] = amp.initialize(model,
[model.optimizer_G, model.optimizer_D],
opt_level='O1')
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
else:
optimizer_G, optimizer_D = model.module.optimizer_G, model.module.optimizer_D
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = epoch_iter % dataset_size
for i, data in enumerate(dataset, start=epoch_iter):
if total_steps % opt.print_freq == print_delta:
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == display_delta
############## Forward Pass ######################
model = model.train()
losses, generated, metrics = model(data['A'],
data['B'],
data['geometry'],
infer=False)
# sum per device losses and metrics
losses = [
torch.mean(x) if not isinstance(x, int) else x for x in losses
]
metric_dict = {k: torch.mean(v) for k, v in metrics.items()}
loss_dict = dict(zip(model.module.loss_names, losses))
# calculate final loss scalar
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + opt.gan_feat_weight * loss_dict.get(
'G_GAN_Feat', 0) + opt.vgg_weight * loss_dict.get('G_VGG', 0)
############### Backward Pass ####################
# update generator weights
optimizer_G.zero_grad()
if opt.fp16:
with amp.scale_loss(loss_G, optimizer_G) as scaled_loss:
scaled_loss.backward()
else:
loss_G.backward()
optimizer_G.step()
# update discriminator weights
optimizer_D.zero_grad()
if opt.fp16:
with amp.scale_loss(loss_D, optimizer_D) as scaled_loss:
scaled_loss.backward()
else:
loss_D.backward()
optimizer_D.step()
############## Display results and errors ##########
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict.items()
}
metrics_ = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in metric_dict.items()
}
t = (time.time() - iter_start_time) / opt.print_freq
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
visualizer.print_current_metrics(epoch, epoch_iter, metrics_,
t)
visualizer.plot_current_metrics(metrics_, total_steps)
#call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"])
### display output images
if save_fake:
if opt.task_type == 'specular':
visuals = OrderedDict([
('albedo', util.tensor2im(data['A'][0])),
('generated',
util.tensor2im_exr(generated.data[0], type=1)),
('GT', util.tensor2im_exr(data['B'][0], type=1))
])
elif opt.task_type == 'low':
visuals = OrderedDict([
('albedo', util.tensor2im(data['A'][0])),
('generated',
util.tensor2im_exr(generated.data[0], type=2)),
('GT', util.tensor2im_exr(data['B'][0], type=2))
])
elif opt.task_type == 'high':
visuals = OrderedDict([
('albedo', util.tensor2im(data['A'][0])),
('generated',
util.tensor2im_exr(generated.data[0], type=3)),
('GT', util.tensor2im_exr(data['B'][0], type=3))
])
visualizer.display_current_results(visuals, epoch, total_steps)
### save latest model
if total_steps % opt.save_latest_freq == save_delta:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch, total_steps))
model.module.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter),
delimiter=',',
fmt='%d')
if epoch_iter >= dataset_size:
break
# end of epoch
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
###########################################################################################
# validation at the end of each epoch
val_start_time = time.time()
metrics_val = []
for _, val_data in enumerate(valset):
model = model.eval()
# model.half()
generated, metrics = model(val_data['A'],
val_data['B'],
val_data['geometry'],
infer=True)
metric_dict = {k: torch.mean(v) for k, v in metrics.items()}
metrics_ = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in metric_dict.items()
}
metrics_val.append(metrics_)
# Print out losses
metrics_val = visualizer.mean4dict(metrics_val)
t = (time.time() - val_start_time) / opt.print_freq
visualizer.print_current_metrics(epoch,
epoch_iter,
metrics_val,
t,
isVal=True)
visualizer.plot_current_metrics(metrics_val, total_steps, isVal=True)
# visualization
if opt.task_type == 'specular':
visuals = OrderedDict([
('albedo', util.tensor2im(val_data['A'][0])),
('generated', util.tensor2im_exr(generated.data[0], type=1)),
('GT', util.tensor2im_exr(val_data['B'][0], type=1))
])
if opt.task_type == 'low':
visuals = OrderedDict([
('albedo', util.tensor2im(val_data['A'][0])),
('generated', util.tensor2im_exr(generated.data[0], type=2)),
('GT', util.tensor2im_exr(val_data['B'][0], type=2))
])
if opt.task_type == 'high':
visuals = OrderedDict([
('albedo', util.tensor2im(val_data['A'][0])),
('generated', util.tensor2im_exr(generated.data[0], type=3)),
('GT', util.tensor2im_exr(val_data['B'][0], type=3))
])
visualizer.display_current_results(visuals, epoch, epoch, isVal=True)
###########################################################################################
### save model for this epoch
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
model.module.save('latest')
model.module.save(epoch)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
### instead of only training the local enhancer, train the entire network after certain iterations
if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
model.module.update_fixed_params()
### linearly decay learning rate after certain iterations
if epoch > opt.niter:
model.module.update_learning_rate()
def train(opt):
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
if opt.continue_train:
if opt.which_epoch == 'latest':
try:
start_epoch, epoch_iter = np.loadtxt(iter_path,
delimiter=',',
dtype=int)
except:
start_epoch, epoch_iter = 1, 0
else:
start_epoch, epoch_iter = int(opt.which_epoch), 0
print('Resuming from epoch %d at iteration %d' %
(start_epoch, epoch_iter))
for update_point in opt.decay_epochs:
if start_epoch < update_point:
break
opt.lr *= opt.decay_gamma
else:
start_epoch, epoch_iter = 0, 0
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = (start_epoch) * dataset_size + epoch_iter
display_delta = total_steps % opt.display_freq
print_delta = total_steps % opt.print_freq
save_delta = total_steps % opt.save_latest_freq
bSize = opt.batchSize
#in case there's no display sample one image from each class to test after every epoch
if opt.display_id == 0:
dataset.dataset.set_sample_mode(True)
dataset.num_workers = 1
for i, data in enumerate(dataset):
if i * opt.batchSize >= opt.numClasses:
break
if i == 0:
sample_data = data
else:
for key, value in data.items():
if torch.is_tensor(data[key]):
sample_data[key] = torch.cat(
(sample_data[key], data[key]), 0)
else:
sample_data[key] = sample_data[key] + data[key]
dataset.num_workers = opt.nThreads
dataset.dataset.set_sample_mode(False)
for epoch in range(start_epoch, opt.epochs):
epoch_start_time = time.time()
if epoch != start_epoch:
epoch_iter = 0
for i, data in enumerate(dataset, start=epoch_iter):
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = (total_steps % opt.display_freq
== display_delta) and (opt.display_id > 0)
############## Network Pass ########################
model.set_inputs(data)
disc_losses = model.update_D()
gen_losses, gen_in, gen_out, rec_out, cyc_out = model.update_G(
infer=save_fake)
loss_dict = dict(gen_losses, **disc_losses)
##################################################
############## Display results and errors ##########
### print out errors
if total_steps % opt.print_freq == print_delta:
errors = {
k: v.item()
if not (isinstance(v, float) or isinstance(v, int)) else v
for k, v in loss_dict.items()
}
t = (time.time() - iter_start_time) / opt.batchSize
visualizer.print_current_errors(epoch + 1, epoch_iter, errors,
t)
if opt.display_id > 0:
visualizer.plot_current_errors(
epoch,
float(epoch_iter) / dataset_size, opt, errors)
### display output images
if save_fake and opt.display_id > 0:
class_a_suffix = ' class {}'.format(data['A_class'][0])
class_b_suffix = ' class {}'.format(data['B_class'][0])
classes = None
visuals = OrderedDict()
visuals_A = OrderedDict([('real image' + class_a_suffix,
util.tensor2im(gen_in.data[0]))])
visuals_B = OrderedDict([('real image' + class_b_suffix,
util.tensor2im(gen_in.data[bSize]))])
A_out_vis = OrderedDict([('synthesized image' + class_b_suffix,
util.tensor2im(gen_out.data[0]))])
B_out_vis = OrderedDict([('synthesized image' + class_a_suffix,
util.tensor2im(gen_out.data[bSize]))
])
if opt.lambda_rec > 0:
A_out_vis.update([('reconstructed image' + class_a_suffix,
util.tensor2im(rec_out.data[0]))])
B_out_vis.update([('reconstructed image' + class_b_suffix,
util.tensor2im(rec_out.data[bSize]))])
if opt.lambda_cyc > 0:
A_out_vis.update([('cycled image' + class_a_suffix,
util.tensor2im(cyc_out.data[0]))])
B_out_vis.update([('cycled image' + class_b_suffix,
util.tensor2im(cyc_out.data[bSize]))])
visuals_A.update(A_out_vis)
visuals_B.update(B_out_vis)
visuals.update(visuals_A)
visuals.update(visuals_B)
ncols = len(visuals_A)
visualizer.display_current_results(visuals, epoch, classes,
ncols)
### save latest model
if total_steps % opt.save_latest_freq == save_delta:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch + 1, total_steps))
model.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter),
delimiter=',',
fmt='%d')
if opt.display_id == 0:
model.eval()
visuals = model.inference(sample_data)
visualizer.save_matrix_image(visuals, 'latest')
model.train()
# end of epoch
iter_end_time = time.time()
print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch + 1, opt.epochs, time.time() - epoch_start_time))
### save model for this epoch
if (epoch + 1) % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch + 1, total_steps))
model.save('latest')
model.save(epoch + 1)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
if opt.display_id == 0:
model.eval()
visuals = model.inference(sample_data)
visualizer.save_matrix_image(visuals, epoch + 1)
model.train()
### multiply learning rate by opt.decay_gamma after certain iterations
if (epoch + 1) in opt.decay_epochs:
model.update_learning_rate()
def train():
opt = TrainOptions().parse()
if opt.distributed:
init_dist()
opt.batchSize = opt.batchSize // len(opt.gpu_ids)
### setup dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
### setup trainer
trainer = Trainer(opt, data_loader)
### setup models
model, flowNet = create_model(opt, trainer.start_epoch)
flow_gt = conf_gt = [None] * 3
ref_idx_fix = torch.zeros([opt.batchSize])
for epoch in tqdm(
range(trainer.start_epoch, opt.niter + opt.niter_decay + 1)):
trainer.start_of_epoch(epoch, model, data_loader)
n_frames_total, n_frames_load = data_loader.dataset.n_frames_total, opt.n_frames_per_gpu
for idx, data in enumerate(tqdm(dataset), start=trainer.epoch_iter):
trainer.start_of_iter()
if not opt.warp_ani:
data.update({
'ani_image': None,
'ani_lmark': None,
'cropped_images': None,
'cropped_lmarks': None
})
if not opt.no_flow_gt:
data_list = [
data['tgt_mask_images'], data['cropped_images'],
data['warping_ref'], data['ani_image']
]
flow_gt, conf_gt = flowNet(data_list, epoch)
data_list = [
data['tgt_label'], data['tgt_image'], data['tgt_template'],
data['cropped_images'], flow_gt, conf_gt
]
data_ref_list = [data['ref_label'], data['ref_image']]
data_prev = [None, None, None]
data_ani = [
data['warping_ref_lmark'], data['warping_ref'],
data['ori_warping_refs'], data['ani_lmark'], data['ani_image']
]
############## Forward Pass ######################
prevs = {"raw_images":[], "synthesized_images":[], \
"prev_warp_images":[], "prev_weights":[], \
"ani_warp_images":[], "ani_weights":[], \
"ref_warp_images":[], "ref_weights":[], \
"ref_flows":[], "prev_flows":[], "ani_flows":[], \
"ani_syn":[]}
for t in range(0, n_frames_total, n_frames_load):
data_list_t = get_data_t(data_list, n_frames_load, t) + data_ref_list + \
get_data_t(data_ani, n_frames_load, t) + data_prev
g_losses, generated, data_prev, ref_idx = model(
data_list_t,
save_images=trainer.save,
mode='generator',
ref_idx_fix=ref_idx_fix)
g_losses = loss_backward(opt, g_losses,
model.module.optimizer_G)
d_losses, _ = model(data_list_t,
mode='discriminator',
ref_idx_fix=ref_idx_fix)
d_losses = loss_backward(opt, d_losses,
model.module.optimizer_D)
# store previous
store_prev(generated, prevs)
loss_dict = dict(
zip(model.module.lossCollector.loss_names,
g_losses + d_losses))
output_data_list = [prevs] + [
data['ref_image']
] + data_ani + data_list + [data['tgt_mask_images']]
if trainer.end_of_iter(loss_dict, output_data_list, model):
break
trainer.end_of_epoch(model)
opt_train = TrainOptions().parse()
opt_val = TrainOptions().parse()
# Random seed
opt_train.use_gpu = len(opt_train.gpu_ids) and torch.cuda.is_available()
if opt_train.manualSeed is None:
opt_train.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt_train.manualSeed)
random.seed(opt_train.manualSeed)
np.random.seed(opt_train.manualSeed)
torch.manual_seed(opt_train.manualSeed)
if opt_train.use_gpu:
torch.cuda.manual_seed_all(opt_train.manualSeed)
data_loader = CreateDataLoader(opt_train)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
opt_val.phase = 'val'
opt_val.nThreads = 1 # test code only supports nThreads = 1
opt_val.batchSize = 1 # test code only supports batchSize = 1
opt_val.serial_batches = True # no shuffle
opt_val.no_flip = True # no flip
opt_val.no_rotate = True # no rotate
if opt_val.valSize == 0:
opt_val.valSize = opt_val.loadSize
opt_val.loadSize = opt_val.valSize
opt_val.fineSize = opt_val.valSize
data_loader_val = CreateDataLoader(opt_val)
help='manual seed')
parser.add_argument(
'--no_lsgan',
action='store_true',
help='do *not* use least square GAN, if false, use vanilla GAN')
parser.add_argument(
'--pool_size',
type=int,
default=62,
help='the size of image buffer that stores previously generated images')
opt = parser.parse_args()
print(opt)
face_dataset = CreateDataLoader(
opt,
csv_fileA='./dataset/celeba/fine_grained_attribute_testA.txt',
root_dirA='./dataset/celeba/testA/',
csv_fileB='./dataset/celeba/fine_grained_attribute_testB.txt',
root_dirB='./dataset/celeba/testB/')
# face_dataset = CreateDataLoader(opt, csv_fileA='./dataset/lfw/fine_grained_attribute_testA.txt',root_dirA='./dataset/lfw/testA/',
# csv_fileB='./dataset/lfw/fine_grained_attribute_testB.txt',root_dirB='./dataset/lfw/testB/')
dataset_size = len(face_dataset)
print('#test images = %d' % dataset_size)
if opt.manualSeed is None:
opt.manualSeed = np.random.randint(1, 10000)
np.random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
torch.cuda.manual_seed_all(opt.manualSeed)
model = create_model(opt)
def train():
opt = TrainOptions().parse()
if opt.debug:
opt.display_freq = 1
opt.print_freq = 1
opt.nThreads = 1
# initialize dataset
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
dataset_size = len(data_loader)
if opt.dataset_mode == 'pose':
print('#training frames = %d' % dataset_size)
else:
print('#training videos = %d' % dataset_size)
# initialize models
modelG, modelD, flowNet = create_model(opt)
visualizer = Visualizer(opt)
iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
# if continue training, recover previous states
if opt.continue_train:
try:
start_epoch, epoch_iter = np.loadtxt(iter_path,
delimiter=',',
dtype=int)
except:
start_epoch, epoch_iter = 1, 0
print('Resuming from epoch %d at iteration %d' %
(start_epoch, epoch_iter))
if start_epoch > opt.niter:
modelG.module.update_learning_rate(start_epoch - 1)
modelD.module.update_learning_rate(start_epoch - 1)
if (opt.n_scales_spatial > 1) and (opt.niter_fix_global != 0) and (
start_epoch > opt.niter_fix_global):
modelG.module.update_fixed_params()
if start_epoch > opt.niter_step:
data_loader.dataset.update_training_batch(
(start_epoch - 1) // opt.niter_step)
modelG.module.update_training_batch(
(start_epoch - 1) // opt.niter_step)
else:
start_epoch, epoch_iter = 1, 0
# set parameters
# number of gpus used for generator for each batch
n_gpus = opt.n_gpus_gen // opt.batchSize
tG, tD = opt.n_frames_G, opt.n_frames_D
tDB = tD * opt.output_nc
s_scales = opt.n_scales_spatial
t_scales = opt.n_scales_temporal
input_nc = 1 if opt.label_nc != 0 else opt.input_nc
output_nc = opt.output_nc
opt.print_freq = lcm(opt.print_freq, opt.batchSize)
total_steps = (start_epoch - 1) * dataset_size + epoch_iter
total_steps = total_steps // opt.print_freq * opt.print_freq
# real training starts here
for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
for idx, data in enumerate(dataset, start=epoch_iter):
if total_steps % opt.print_freq == 0:
iter_start_time = time.time()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
# whether to collect output images
save_fake = total_steps % opt.display_freq == 0
# n_frames_total = n_frames_load * n_loadings + tG - 1
_, n_frames_total, height, width = data['B'].size()
n_frames_total = n_frames_total // opt.output_nc
# number of total frames loaded into GPU at a time for each batch
n_frames_load = opt.max_frames_per_gpu * n_gpus
n_frames_load = min(n_frames_load, n_frames_total - tG + 1)
# number of loaded frames plus previous frames
t_len = n_frames_load + tG - 1
# the last generated frame from previous training batch (which becomes input to the next batch)
fake_B_last = None
# all real/generated frames so far
real_B_all, fake_B_all, flow_ref_all, conf_ref_all = None, None, None, None
if opt.sparse_D:
real_B_all, fake_B_all, flow_ref_all, conf_ref_all = [
None
] * t_scales, [None] * t_scales, [None] * t_scales, [
None
] * t_scales
# temporally subsampled frames
real_B_skipped, fake_B_skipped = [None] * t_scales, [None
] * t_scales
flow_ref_skipped, conf_ref_skipped = [None] * t_scales, [
None
] * t_scales # temporally subsampled flows
for i in range(0, n_frames_total - t_len + 1, n_frames_load):
# 5D tensor: batchSize, # of frames, # of channels, height, width
input_A = Variable(
data['A'][:, i * input_nc:(i + t_len) * input_nc,
...]).view(-1, t_len, input_nc, height, width)
input_B = Variable(
data['B'][:, i * output_nc:(i + t_len) * output_nc,
...]).view(-1, t_len, output_nc, height, width)
inst_A = Variable(data['inst'][:, i:i + t_len, ...]).view(
-1, t_len, 1, height,
width) if len(data['inst'].size()) > 2 else None
###################################### Forward Pass ##########################
# generator
fake_B, fake_B_raw, flow, weight, real_A, real_Bp, fake_B_last = modelG(
input_A, input_B, inst_A, fake_B_last)
if i == 0:
# the first generated image in this sequence
fake_B_first = fake_B[0, 0]
# the collection of previous and current real frames
real_B_prev, real_B = real_Bp[:, :-1], real_Bp[:, 1:]
# discriminator
# individual frame discriminator
# reference flows and confidences
flow_ref, conf_ref = flowNet(real_B, real_B_prev)
fake_B_prev = real_B_prev[:, 0:
1] if fake_B_last is None else fake_B_last[
0][:, -1:]
if fake_B.size()[1] > 1:
fake_B_prev = torch.cat(
[fake_B_prev, fake_B[:, :-1].detach()], dim=1)
losses = modelD(
0,
reshape([
real_B, fake_B, fake_B_raw, real_A, real_B_prev,
fake_B_prev, flow, weight, flow_ref, conf_ref
]))
losses = [
torch.mean(x) if x is not None else 0 for x in losses
]
loss_dict = dict(zip(modelD.module.loss_names, losses))
# temporal discriminator
loss_dict_T = []
# get skipped frames for each temporal scale
if t_scales > 0:
if opt.sparse_D:
real_B_all, real_B_skipped = get_skipped_frames_sparse(
real_B_all, real_B, t_scales, tD, n_frames_load, i)
fake_B_all, fake_B_skipped = get_skipped_frames_sparse(
fake_B_all, fake_B, t_scales, tD, n_frames_load, i)
flow_ref_all, flow_ref_skipped = get_skipped_frames_sparse(
flow_ref_all,
flow_ref,
t_scales,
tD,
n_frames_load,
i,
is_flow=True)
conf_ref_all, conf_ref_skipped = get_skipped_frames_sparse(
conf_ref_all,
conf_ref,
t_scales,
tD,
n_frames_load,
i,
is_flow=True)
else:
real_B_all, real_B_skipped = get_skipped_frames(
real_B_all, real_B, t_scales, tD)
fake_B_all, fake_B_skipped = get_skipped_frames(
fake_B_all, fake_B, t_scales, tD)
flow_ref_all, conf_ref_all, flow_ref_skipped, conf_ref_skipped = get_skipped_flows(
flowNet, flow_ref_all, conf_ref_all,
real_B_skipped, flow_ref, conf_ref, t_scales, tD)
# run discriminator for each temporal scale
for s in range(t_scales):
if real_B_skipped[s] is not None:
losses = modelD(s + 1, [
real_B_skipped[s], fake_B_skipped[s],
flow_ref_skipped[s], conf_ref_skipped[s]
])
losses = [
torch.mean(x) if not isinstance(x, int) else x
for x in losses
]
loss_dict_T.append(
dict(zip(modelD.module.loss_names_T, losses)))
# collect losses
loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
loss_G = loss_dict['G_GAN'] + \
loss_dict['G_GAN_Feat'] + loss_dict['G_VGG']
loss_G += loss_dict['G_Warp'] + loss_dict['F_Flow'] + \
loss_dict['F_Warp'] + loss_dict['W']
if opt.add_face_disc:
loss_G += loss_dict['G_f_GAN'] + loss_dict['G_f_GAN_Feat']
loss_D += (loss_dict['D_f_fake'] +
loss_dict['D_f_real']) * 0.5
# collect temporal losses
loss_D_T = []
t_scales_act = min(t_scales, len(loss_dict_T))
for s in range(t_scales_act):
loss_G += loss_dict_T[s]['G_T_GAN'] + \
loss_dict_T[s]['G_T_GAN_Feat'] + \
loss_dict_T[s]['G_T_Warp']
loss_D_T.append((loss_dict_T[s]['D_T_fake'] +
loss_dict_T[s]['D_T_real']) * 0.5)
###################################### Backward Pass #################################
optimizer_G = modelG.module.optimizer_G
optimizer_D = modelD.module.optimizer_D
# update generator weights
optimizer_G.zero_grad()
loss_G.backward()
optimizer_G.step()
# update discriminator weights
# individual frame discriminator
optimizer_D.zero_grad()
loss_D.backward()
optimizer_D.step()
# temporal discriminator
for s in range(t_scales_act):
optimizer_D_T = getattr(modelD.module,
'optimizer_D_T' + str(s))
optimizer_D_T.zero_grad()
loss_D_T[s].backward()
optimizer_D_T.step()
if opt.debug:
call([
"nvidia-smi", "--format=csv",
"--query-gpu=memory.used,memory.free"
])
############## Display results and errors ##########
# print out errors
if total_steps % opt.print_freq == 0:
t = (time.time() - iter_start_time) / opt.print_freq
errors = {
k: v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict.items()
}
for s in range(len(loss_dict_T)):
errors.update({
k + str(s):
v.data.item() if not isinstance(v, int) else v
for k, v in loss_dict_T[s].items()
})
visualizer.print_current_errors(epoch, epoch_iter, errors, t)
visualizer.plot_current_errors(errors, total_steps)
# display output images
if save_fake:
if opt.label_nc != 0:
input_image = util.tensor2label(real_A[0, -1],
opt.label_nc)
elif opt.dataset_mode == 'pose':
input_image = util.tensor2im(real_A[0, -1, :3])
if real_A.size()[2] == 6:
input_image2 = util.tensor2im(real_A[0, -1, 3:])
input_image[input_image2 != 0] = input_image2[
input_image2 != 0]
else:
c = 3 if opt.input_nc == 3 else 1
input_image = util.tensor2im(real_A[0, -1, :c],
normalize=False)
if opt.use_instance:
edges = util.tensor2im(real_A[0, -1, -1:, ...],
normalize=False)
input_image += edges[:, :, np.newaxis]
if opt.add_face_disc:
ys, ye, xs, xe = modelD.module.get_face_region(real_A[0,
-1:])
if ys is not None:
input_image[ys, xs:xe, :] = input_image[
ye, xs:xe, :] = input_image[
ys:ye, xs, :] = input_image[ys:ye, xe, :] = 255
visual_list = [
('input_image', input_image),
('fake_image', util.tensor2im(fake_B[0, -1])),
('fake_first_image', util.tensor2im(fake_B_first)),
('fake_raw_image', util.tensor2im(fake_B_raw[0, -1])),
('real_image', util.tensor2im(real_B[0, -1])),
('flow_ref', util.tensor2flow(flow_ref[0, -1])),
('conf_ref',
util.tensor2im(conf_ref[0, -1], normalize=False))
]
if flow is not None:
visual_list += [('flow', util.tensor2flow(flow[0, -1])),
('weight',
util.tensor2im(weight[0, -1],
normalize=False))]
visuals = OrderedDict(visual_list)
visualizer.display_current_results(visuals, epoch, total_steps)
# save latest model
if total_steps % opt.save_latest_freq == 0:
visualizer.vis_print(
'saving the latest model (epoch %d, total_steps %d)' %
(epoch, total_steps))
modelG.module.save('latest')
modelD.module.save('latest')
np.savetxt(iter_path, (epoch, epoch_iter),
delimiter=',',
fmt='%d')
if epoch_iter > dataset_size - opt.batchSize:
epoch_iter = 0
break
# end of epoch
iter_end_time = time.time()
visualizer.vis_print('End of epoch %d / %d \t Time Taken: %d sec' %
(epoch, opt.niter + opt.niter_decay,
time.time() - epoch_start_time))
# save model for this epoch
if epoch % opt.save_epoch_freq == 0:
visualizer.vis_print(
'saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
modelG.module.save('latest')
modelD.module.save('latest')
modelG.module.save(epoch)
modelD.module.save(epoch)
np.savetxt(iter_path, (epoch + 1, 0), delimiter=',', fmt='%d')
# linearly decay learning rate after certain iterations
if epoch > opt.niter:
modelG.module.update_learning_rate(epoch)
modelD.module.update_learning_rate(epoch)
# gradually grow training sequence length
if (epoch % opt.niter_step) == 0:
data_loader.dataset.update_training_batch(epoch // opt.niter_step)
modelG.module.update_training_batch(epoch // opt.niter_step)
# finetune all scales
if (opt.n_scales_spatial > 1) and (opt.niter_fix_global != 0) and (
epoch == opt.niter_fix_global):
modelG.module.update_fixed_params()
import time
from options.train_options import TrainOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
import pudb
opt = TrainOptions().parse()
data_loader = CreateDataLoader(opt) # inits CustomDatasetDataLoader
dataset = data_loader.load_data() # returns CustomDatasetDataLoader obj
dataset_size = len(data_loader)
print('#training images = %d' % dataset_size)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = 0
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
epoch_start_time = time.time()
iter_data_time = time.time()
epoch_iter = 0
skip_count = 0
for i, data in enumerate(dataset):
iter_start_time = time.time()
if total_steps % opt.print_freq == 0:
t_data = iter_start_time - iter_data_time
visualizer.reset()
total_steps += opt.batchSize
epoch_iter += opt.batchSize
from options.encoder_decoder_options_v2 import TrainEncoderDecoderOptions_V2
from misc.visualizer import GANVisualizer_V3
import util.io as io
import os
import sys
import time
import numpy as np
from collections import OrderedDict
opt = TrainEncoderDecoderOptions_V2().parse()
# create model
model = EncoderDecoderFramework_DFN()
model.initialize(opt)
# create data loader
train_loader = CreateDataLoader(opt, split = 'train')
val_loader = CreateDataLoader(opt, split = 'test')
# create visualizer
visualizer = GANVisualizer_V3(opt)
total_steps = 0
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
model.update_learning_rate()
for i, data in enumerate(train_loader):
total_steps += 1
model.set_input(data)
model.forward()
model.optimize_parameters()
if total_steps % opt.display_freq == 0:
import time
from options.test_options import TestOptions
from options.train_options import TrainOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
# train
opt_train = TrainOptions().parse()
data_loader_train = CreateDataLoader(opt_train)
dataset_train = data_loader_train.load_data()
dataset_size_train = len(data_loader_train)
print('#training images = %d' % dataset_size_train)
# test
opt_test = TestOptions().parse()
opt_test.nThreads = 1 # test code only supports nThreads = 1
opt_test.batchSize = 1 # test code only supports batchSize = 1
opt_test.serial_batches = False # no shuffle
opt_test.no_flip = True # no flip
opt_test.how_many = 100
data_loader_test = CreateDataLoader(opt_test)
dataset_test = data_loader_test.load_data()
dataset_size_test = len(data_loader_test)
print('#test images = %d' % dataset_size_test)
model = create_model(opt_train)
visualizer = Visualizer(opt_train)
total_steps = 0
for epoch in range(opt_train.epoch_count,
loss_fns = [GramMSELoss()] * len(loss_layer)
if torch.cuda.is_available():
loss_fns = [loss_fn.cuda() for loss_fn in loss_fns]
style_weights = [1e3 / n**2 for n in [64, 128, 256, 512]]
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()), 0.0001,
[0.9, 0.999])
optimizer_D = torch.optim.Adam(
filter(lambda p: p.requires_grad, netD.parameters()), 0.0002, [0.9, 0.999])
print('# generator parameters:',
sum(param.numel() for param in model.parameters()))
print('# discriminator parameters:',
sum(param.numel() for param in netD.parameters()))
print()
for i in range(opt.start_training_step, 4):
opt.nEpochs = training_settings[i - 1]['nEpochs']
opt.lr = training_settings[i - 1]['lr']
opt.step = training_settings[i - 1]['step']
opt.lr_decay = training_settings[i - 1]['lr_decay']
opt.lambda_db = training_settings[i - 1]['lambda_db']
opt.gated = training_settings[i - 1]['gated']
print(opt)
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
lr = adjust_learning_rate(epoch - 1)
trainloader = CreateDataLoader(opt)
train(trainloader, model, criterion, optimizer, epoch, lr)
if epoch % 5 == 0:
checkpoint(i, epoch)
# load options
opt = TestMMGANOptions_V3().parse()
opt.batch_size = batch_size
org_opt = io.load_json(os.path.join('checkpoints', opt.id, 'train_opt.json'))
preserved_opt = {'id', 'gpu_ids', 'batch_size', 'is_train', 'dataset_mode'}
for k, v in org_opt.iteritems():
if k in opt and (k not in preserved_opt):
setattr(opt, k, v)
# create model
model = MultimodalDesignerGAN_V3()
model.initialize(opt)
for name, net in model.modules.iteritems():
for p in net.parameters():
p.requires_grad = False
# create data loader
val_loader_iter = iter(CreateDataLoader(opt, split='vis'))
# create visualizer
visualizer = GANVisualizer_V3(opt)
for i in range(num_batch):
print('[%s] test generation: %d / %d' % (opt.id, i + 1, num_batch))
data = val_loader_iter.next()
imgs_edge = data['img_edge'][0::7].cpu().clone()
imgs_color = data['img_color'][0:7].cpu().clone()
s_id = data['id'][0][0]
model.set_input(data)
model.test('normal')
vis_dir = 'vis_gen'
import time
import os
from options.test_options import TestOptions
from data.data_loader import CreateDataLoader
from models.models import create_model
from util.visualizer import Visualizer
from util import html
opt = TestOptions().parse()
opt.nThreads = 1 # test code only supports nThreads = 1
opt.batchSize = 1 # test code only supports batchSize = 1
opt.serial_batches = True # no shuffle
opt.no_flip = True # no flip
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
model = create_model(opt)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(
web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
for i, data in enumerate(dataset):
if i >= opt.how_many:
break
model.set_input(data)
model.test()
visuals = model.get_current_visuals_test()
import torch.backends.cudnn as cudnn
from torchsummary import summary
from tqdm import tqdm
from data.data_loader import CreateDataLoader
from models.models import create_model
from options.train_options import TrainOptions
cudnn.benchmark = True
opt = TrainOptions().parse()
data_loader = CreateDataLoader(
opt, {
'mode':
'Train',
'manifestFn':
'/home/caspardu/data/LipReadProject/LipNetData/manifestFiles/wellDone_train.list',
'labelFn':
'/home/caspardu/data/LipReadProject/LipNetData/manifestFiles/label.txt'
})
testDataLoader = CreateDataLoader(
opt, {
'mode':
'Test',
'manifestFn':
'/home/caspardu/data/LipReadProject/LipNetData/manifestFiles/wellDone_test.list',
'labelFn':
'/home/caspardu/data/LipReadProject/LipNetData/manifestFiles/label.txt'
})
dataset = data_loader.load_data()
