import copy import pdb 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 #copy opt2 = copy.copy(opt) opt2.dataroot = opt.dataroot_more data_loader = CreateDataLoader(opt) data_loader2 = CreateDataLoader(opt2) dataset_size = min(len(data_loader), len(data_loader2)) dataset = data_loader.load_data() dataset2 = data_loader2.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 loader = enumerate(dataset) loader2 = enumerate(dataset2) for steps in range(dataset_size): i, data = next(loader) i, data2 = next(loader2)
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()
from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import ntpath import os from util import util import shutil import numpy as np #from util.visualizer import Visualizer opt = TrainOptions().parse() # Load dataset_train data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) # Load dataset_test opt.batchSize = 1 opt.phase = 'test' data_loader_test = CreateDataLoader(opt) dataset_test = data_loader_test.load_data() dataset_size = len(data_loader_test) print('#testing images = %d' % dataset_size) # Create or clear dir for saving generated samples if os.path.exists(opt.testing_path): shutil.rmtree(opt.testing_path) os.makedirs(opt.testing_path)
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 n_gpus = opt.n_gpus_gen // opt.batchSize # number of gpus used for generator for each batch 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, height, width = data['B'].size( ) # n_frames_total = n_frames_load * n_loadings + tG - 1 n_frames_total = n_frames_total // opt.output_nc n_frames_load = opt.max_frames_per_gpu * n_gpus # number of total frames loaded into GPU at a time for each batch n_frames_load = min(n_frames_load, n_frames_total - tG + 1) t_len = n_frames_load + tG - 1 # number of loaded frames plus previous frames fake_B_last = None # the last generated frame from previous training batch (which becomes input to the next batch) real_B_all, fake_B_all, flow_ref_all, conf_ref_all = None, None, None, None # all real/generated frames so far 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 real_B_skipped, fake_B_skipped = [None] * t_scales, [ None ] * t_scales # temporally subsampled frames 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: fake_B_first = fake_B[ 0, 0] # the first generated image in this sequence real_B_prev, real_B = real_Bp[:, : -1], real_Bp[:, 1:] # the collection of previous and current real frames ####### discriminator ### individual frame discriminator flow_ref, conf_ref = flowNet( real_B, real_B_prev) # reference flows and confidences 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=True) if opt.use_instance: edges = util.tensor2im(real_A[0, -1, -1:, ...], normalize=True) 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()
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()
from options.test_options import TestOptions from data.data_loader import CreateDataLoader from models.create_model import create_model from utils import visualizer opt = TestOptions().parse() ### set dataloader ### print('### prepare DataLoader ###') data_loader = CreateDataLoader(opt) test_loader = data_loader.load_data() print('test images : {}'.format(len(data_loader))) print('numof_iteration : {}'.format(len(test_loader))) ### define model ### model = create_model(opt) model.gen.eval() print('### start test ! ###') for iter, data in enumerate(test_loader): model.set_variables(data) fake_image = model.forward() fake_image = (fake_image + 1.0) / 2.0 visualizer.save_test_images(opt, iter, data['label'], fake_image, nrow=7)
else: break avg_loss = sum(losses) / len(losses) if opt.tensorboard: writer.add_scalar('data/val_loss', avg_loss, index) print('val loss: %.3f' % avg_loss) return avg_loss #parse arguments opt = TrainOptions().parse() opt.device = torch.device("cuda") #construct data loader data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training clips = %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 clips = %d' % dataset_size_val) opt.mode = 'train' #set it back if opt.tensorboard: from tensorboardX import SummaryWriter
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']) cer = sum(totalCer) * 100 / len(totalCer) wer = sum(totalWer) * 100 / len(totalWer) print('############################')
type=int, default=100, help='frequency of showing training results on console') parser.add_argument('--ckpt_path', default='') opt = parser.parse_args() print(opt) train_visual = Visualizer(opt.train_display_id, 'train', 5) val_visual = Visualizer(opt.val_display_id, 'val', 5) if not os.path.exists(opt.ckpt_path): os.makedirs(opt.ckpt_path) data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) opt.phase = 'val' #opt.pahse = 'val' val_data_loader = CreateDataLoader(opt) val_dataset = val_data_loader.load_data() # val_loader = DataLoader(val_set, batch_size=1, num_workers=12, pin_memory=True) ## define models # 1---256x256 stage # 2---128x128 stage # 3---64x64 stage GA = nets.define_G(input_nc=3, output_nc=3,
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 frames = %d' % dataset_size) ### initialize models models = create_model_cloth(opt) ClothWarper, ClothWarperLoss, flowNet, optimizer = create_optimizer_cloth(opt, models) ### set parameters n_gpus, tG, input_nc_1, input_nc_2, input_nc_3, start_epoch, epoch_iter, print_freq, total_steps, iter_path, tD, t_scales = init_params(opt, ClothWarper, data_loader) 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_cloth(data, n_gpus, tG) flow_total_prev_last, frames_all = data_loader.dataset.init_data_cloth(t_scales) for i in range(0, n_frames_total, n_frames_load): is_first_frame = flow_total_prev_last is None input_TParsing, input_TFG, input_SParsing, input_SFG, input_SFG_full = data_loader.dataset.prepare_data_cloth(data, i) ###################################### Forward Pass ########################## ####### C2F-FWN fg_tps, fg_dense, lo_tps, lo_dense, flow_tps, flow_dense, flow_totalp, real_input_1, real_input_2, real_SFG, real_SFG_fullp, flow_total_last = ClothWarper(input_TParsing, input_TFG, input_SParsing, input_SFG, input_SFG_full, flow_total_prev_last) real_SLO = real_input_2[:, :, -opt.label_nc_2:] ####### compute losses ### individual frame losses and FTC loss with l=1 real_SFG_full_prev, real_SFG_full = real_SFG_fullp[:, :-1], real_SFG_fullp[:, 1:] # the collection of previous and current real frames flow_optical_ref, conf_optical_ref = flowNet(real_SFG_full, real_SFG_full_prev) # reference flows and confidences flow_total_prev, flow_total = flow_totalp[:, :-1], flow_totalp[:, 1:] if is_first_frame: flow_total_prev = flow_total_prev[:, 1:] flow_total_prev_last = flow_total_last losses, flows_sampled_0 = ClothWarperLoss(0, reshape([real_SFG, real_SLO, fg_tps, fg_dense, lo_tps, lo_dense, flow_tps, flow_dense, flow_total, flow_total_prev, flow_optical_ref, conf_optical_ref]), is_first_frame) losses = [ torch.mean(x) if x is not None else 0 for x in losses ] loss_dict = dict(zip(ClothWarperLoss.module.loss_names, losses)) ### FTC losses with l=3,9 # get skipped frames for each temporal scale frames_all, frames_skipped = ClothWarperLoss.module.get_all_skipped_frames(frames_all, \ real_SFG_full, flow_total, flow_optical_ref, conf_optical_ref, real_SLO, t_scales, tD, n_frames_load, i, flowNet) # compute losses for l=3,9 loss_dict_T = [] for s in range(1, t_scales): if frames_skipped[0][s] is not None and not opt.tps_only: losses, flows_sampled_1 = ClothWarperLoss(s+1, [frame_skipped[s] for frame_skipped in frames_skipped], False) losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ] loss_dict_T.append(dict(zip(ClothWarperLoss.module.loss_names_T, losses))) # collect losses loss, _ = ClothWarperLoss.module.get_losses(loss_dict, loss_dict_T, t_scales-1) ###################################### Backward Pass ################################# # update generator weights loss_backward(opt, loss, optimizer) if i == 0: fg_dense_first = fg_dense[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()}) loss_names_vis = ClothWarperLoss.module.loss_names.copy() {loss_names_vis.append(ClothWarperLoss.module.loss_names_T[0]+str(idx)) for idx in range(len(loss_dict_T))} visualizer.print_current_errors_new(epoch, epoch_iter, errors, loss_names_vis, t) visualizer.plot_current_errors(errors, total_steps) ### display output images if save_fake: visuals = util.save_all_tensors_cloth(opt, real_input_1, real_input_2, fg_tps, fg_dense, lo_tps, lo_dense, fg_dense_first, real_SFG, real_SFG_full, flow_tps, flow_dense, flow_total) visualizer.display_current_results(visuals, epoch, total_steps) ### save latest model save_models_cloth(opt, epoch, epoch_iter, total_steps, visualizer, iter_path, ClothWarper) 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_cloth(opt, epoch, epoch_iter, total_steps, visualizer, iter_path, ClothWarper, end_of_epoch=True) update_models_cloth(opt, epoch, ClothWarper, data_loader)
def train_target(param_dir, if_train, loadSize, continue_train = False): # import pix2pixHD module pix2pixhd_dir = Path('../src/pix2pixHD/') # should download github code 'pix2pixHD' import sys sys.path.append(str(pix2pixhd_dir)) #get_ipython().run_line_magic('load_ext', 'autoreload') #get_ipython().run_line_magic('autoreload', '2') from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model import util.util as util from util.visualizer import Visualizer from util import html with open(param_dir, mode='rb') as f: opt = pickle.load(f) opt.loadSize = loadSize opt.fineSize = loadSize #opt.fine_size = loadSize #opt.loadsize = loadSize opt.continue_train = continue_train print('opt parameters: ', 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) model = create_model(opt) visualizer = Visualizer(opt) if if_train: print('#training images = %d' % dataset_size) if opt.continue_train: print ('Resumming training ...') else: print ('Starting new training ...') 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 for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): # (1,20+20+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.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'] + loss_dict.get('G_GAN_Feat',0) + loss_dict.get('G_VGG',0) ############### Backward Pass #################### # update generator weights model.module.optimizer_G.zero_grad() loss_G.backward() model.module.optimizer_G.step() # update discriminator weights model.module.optimizer_D.zero_grad() loss_D.backward() model.module.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: 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 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: # opt.save_epoch_freq == 10 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() torch.cuda.empty_cache() else: # 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)) for data in tqdm(dataset): minibatch = 1 generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)), ('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path) webpage.save() torch.cuda.empty_cache()
def prepare_face_enhancer_data(target_dir, run_name): face_sync_dir = os.path.join(target_dir, 'face') os.makedirs(face_sync_dir, exist_ok=True) test_sync_dir = os.path.join(face_sync_dir, 'test_sync') os.makedirs(test_sync_dir, exist_ok=True) test_real_dir = os.path.join(face_sync_dir, 'test_real') os.makedirs(test_real_dir, exist_ok=True) test_img = os.path.join(target_dir, 'test_img') os.makedirs(test_img, exist_ok=True) test_label = os.path.join(target_dir, 'test_label') os.makedirs(test_label, exist_ok=True) transfer_face_sync_dir = os.path.join(target_dir, 'face_transfer') os.makedirs(transfer_face_sync_dir, exist_ok=True) transfer_test_sync_dir = os.path.join(transfer_face_sync_dir, 'test_sync') os.makedirs(transfer_test_sync_dir, exist_ok=True) transfer_test_real_dir = os.path.join(transfer_face_sync_dir, 'test_real') os.makedirs(transfer_test_real_dir, exist_ok=True) train_dir = os.path.join(target_dir, 'train', 'train_img') label_dir = os.path.join(target_dir, 'train', 'train_label') print('Prepare test_real....') for img_file in tqdm(sorted(os.listdir(train_dir))): img_idx = int(img_file.split('.')[0]) img = cv2.imread(os.path.join(train_dir, '{:05}.png'.format(img_idx))) label = cv2.imread(os.path.join(label_dir, '{:05}.png'.format(img_idx))) cv2.imwrite(os.path.join(test_real_dir, '{:05}.png'.format(img_idx)), img) cv2.imwrite( os.path.join(transfer_test_real_dir, '{:05}.png'.format(img_idx)), img) cv2.imwrite(os.path.join(test_img, '{:05}.png'.format(img_idx)), img) cv2.imwrite(os.path.join(test_label, '{:05}.png'.format(img_idx)), label) print('Prepare test_sync....') import src.config.test_opt as opt if device == torch.device('cpu'): opt.gpu_ids = [] else: os.environ['CUDA_VISIBLE_DEVICES'] = "0" opt.checkpoints_dir = os.path.join(dir_name, '../../checkpoints/') opt.dataroot = target_dir opt.name = run_name opt.nThreads = 0 opt.results_dir = os.path.join(dir_name, '../../face_enhancer_results/') iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() visualizer = Visualizer(opt) 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)) model = create_model(opt) for data in tqdm(dataset): minibatch = 1 generated = model.inference(data['label'], data['inst']) visuals = OrderedDict([('synthesized_image', util.tensor2im(generated.data[0]))]) img_path = data['path'] visualizer.save_images(webpage, visuals, img_path) webpage.save() torch.cuda.empty_cache() print(f'Copy the synthesized images in {test_sync_dir}...') synthesized_image_dir = os.path.join(dir_name, '../../face_enhancer_results', run_name, 'test_latest/images/') img_list = [ f for f in os.listdir(synthesized_image_dir) if f.endswith('synthesized_image.jpg') ] for img_file in tqdm(sorted(img_list)): img_idx = int(img_file.split('_')[0]) img = cv2.imread( os.path.join(synthesized_image_dir, '{:05}_synthesized_image.jpg'.format(img_idx))) cv2.imwrite(os.path.join(test_sync_dir, '{:05}.png'.format(img_idx)), img) print('Copy transfer_test_sync') previous_run_img_dir = os.path.join(dir_name, '../../results', run_name, 'test_latest/images/') img_list = [ f for f in os.listdir(previous_run_img_dir) if f.endswith('synthesized_image.jpg') ] for img_file in tqdm(sorted(img_list)): img_idx = int(img_file.split('_')[0]) img = cv2.imread( os.path.join(previous_run_img_dir, '{:05}_synthesized_image.jpg'.format(img_idx))) cv2.imwrite( os.path.join(transfer_test_sync_dir, '{:05}.png'.format(img_idx)), img)
dtype=int) except: start_epoch, epoch_iter = 1, 0 try: best_iou = np.loadtxt(ioupath_path, dtype=float) except: best_iou = 0. print('Resuming from epoch %d at iteration %d, previous best IoU %f' % (start_epoch, epoch_iter, best_iou)) else: start_epoch, epoch_iter = 1, 0 best_iou = 0. data_loader = CreateDataLoader(opt) dataset, dataset_val = data_loader.load_data() dataset_size = len(dataset) print('#training images = %d' % dataset_size) model = create_model(opt, dataset.dataset) # print (model) visualizer = Visualizer(opt) total_steps = (start_epoch - 1) * dataset_size + epoch_iter for epoch in range(start_epoch, opt.nepochs): epoch_start_time = time.time() if epoch != start_epoch: epoch_iter = epoch_iter % dataset_size model.model.train() for i, data in enumerate(dataset, start=epoch_iter): iter_start_time = time.time()
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()
def test(opt): 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() dataset_size = len(data_loader) print('#test batches = %d' % (int(dataset_size / len(opt.sort_order)))) visualizer = Visualizer(opt) model = create_model(opt) model.eval() # create webpage if opt.random_seed != -1: exp_dir = '%s_%s_seed%s' % (opt.phase, opt.which_epoch, str(opt.random_seed)) else: exp_dir = '%s_%s' % (opt.phase, opt.which_epoch) web_dir = os.path.join(opt.results_dir, opt.name, exp_dir) if opt.traverse or opt.deploy: if opt.traverse: out_dirname = 'traversal' else: out_dirname = 'deploy' output_dir = os.path.join(web_dir, out_dirname) if not os.path.isdir(output_dir): os.makedirs(output_dir) for image_path in opt.image_path_list: print(image_path) data = dataset.dataset.get_item_from_path(image_path) visuals = model.inference(data) if opt.traverse and opt.make_video: out_path = os.path.join( output_dir, os.path.splitext(os.path.basename(image_path))[0] + '.mp4') visualizer.make_video(visuals, out_path) elif opt.traverse or (opt.deploy and opt.full_progression): if opt.traverse and opt.compare_to_trained_outputs: out_path = os.path.join( output_dir, os.path.splitext(os.path.basename(image_path))[0] + '_compare_to_{}_jump_{}.png'.format( opt.compare_to_trained_class, opt.trained_class_jump)) else: out_path = os.path.join( output_dir, os.path.splitext(os.path.basename(image_path))[0] + '.png') visualizer.save_row_image(visuals, out_path, traverse=opt.traverse) else: out_path = os.path.join(output_dir, os.path.basename(image_path[:-4])) visualizer.save_images_deploy(visuals, out_path) else: 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 visuals = model.inference(data) img_path = data['Paths'] rem_ind = [] for i, path in enumerate(img_path): if path != '': print('process image... %s' % path) else: rem_ind += [i] for ind in reversed(rem_ind): del img_path[ind] visualizer.save_images(webpage, visuals, img_path) webpage.save()
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 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 epoch_iter = 0 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'] + 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()
from pdb import set_trace as st from util import html opt = TrainOptions().parse() opt.manualSeed = random.randint(1, 10000) # fix seed print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) avi2pngs(opt) if opt.file_list != None: opt.dataroot = opt.dataroot + '/split' data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training samples = %d' % dataset_size) train_batch_size = opt.batchSize model = create_model(opt) visualizer = Visualizer(opt) 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
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, data_loader) 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_prev_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_prev_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 fake_B_prev = modelG.module.compute_fake_B_prev( real_B_prev, fake_B_prev_last, fake_B) fake_B_prev_last = fake_B_last 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 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()
def main(): # 입력 X, return X iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') # 반복 경로 받아오기 data_loader = CreateDataLoader(opt) # option에 해당하는 data_loader 생성 dataset = data_loader.load_data() # dataset을 data_loader로부터 받아온다. dataset_size = len(data_loader) # dataset의 사이즈를 지정 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 # delta 값들 지정 model = create_model(opt) # model = model.cuda() visualizer = Visualizer(opt) # 현재 option에 해당하는 훈련 과정 출력 for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1): # 총 40번 반복 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: 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') 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()
import time from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.models import create_model opt = TrainOptions().parse() data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() dataset_size = len(data_loader) print('#training images = %d' % dataset_size) model = create_model(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 for i, data in enumerate(dataset): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter += opt.batchSize model.set_input(data) model.optimize_parameters() if total_steps % opt.print_freq == 0: errors = model.get_current_errors() t = (time.time() - iter_start_time) / opt.batchSize
import time from options.train_options import TrainOptions opt = TrainOptions().parse() from data.data_loader import CreateDataLoader from models.models import create_model from util.logger import Logger data_loader = CreateDataLoader(opt) dataset = data_loader.load_data( ) # dataset is actually a torch.utils.data.DataLoader object dataset_size = len(data_loader) print('#training images = %d' % dataset_size) logger = Logger(opt) model = create_model(opt) total_steps = 0 for epoch in range(1, opt.niter + opt.niter_decay + 1): epoch_start_time = time.time() for i, data in enumerate(dataset): iter_start_time = time.time() total_steps += opt.batchSize epoch_iter = total_steps - dataset_size * ( epoch - 1) # iter index in current epoch model.set_input(data) model.optimize_parameters()
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() img_path = model.get_image_paths() print('process image... %s' % img_path) visualizer.save_images(webpage, visuals, img_path)
from options.train_options import TrainOptions from data.data_loader import CreateDataLoader from models.create_model import create_model from utils import visualizer opt = TrainOptions().parse() ### set dataloader ### print('### prepare DataLoader ###') data_loader = CreateDataLoader(opt) train_loader = data_loader.load_data() print('training images : {}'.format(len(data_loader))) print('numof_iteration : {}'.format(len(train_loader))) ### define model ### model = create_model(opt) ### training loop ### print('### start training ! ###') for epoch in range(opt.epoch+opt.epoch_decay+1): for iter, data in enumerate(train_loader): model.set_variables(data) model.optimize_parameters() if iter % opt.print_iter_freq == 0: losses = model.get_current_losses() visualizer.print_current_losses(epoch, iter, losses) image = (data['image'] + 1.0) / 2.0