parser = TestFlowRegressionOptions() opt = parser.parse(display=False) parser.save() print('load training options.') train_opt = io.load_json(os.path.join('checkpoints', opt.id, 'train_opt.json')) preserved_opt = {'gpu_ids', 'is_train', 'batch_size', 'which_epoch', 'debug'} for k, v in train_opt.iteritems(): if k in opt and (k not in preserved_opt): setattr(opt, k, v) # create model model = FlowRegressionModel() model.initialize(opt) # create data loader val_loader = CreateDataLoader(opt, split='test') # create visualizer visualizer = Visualizer(opt) # test loss_buffer = LossBuffer(size=len(val_loader)) model.output = {} model.eval() for i, data in enumerate(tqdm.tqdm(val_loader, desc='Test')): model.set_input(data) model.test(compute_loss=True) loss_buffer.add(model.get_current_errors()) errors = loss_buffer.get_errors() info = OrderedDict([('model_id', opt.id), ('epoch', opt.which_epoch)]) log_str = visualizer.log(info, errors, log_in_file=False) print(log_str)
for i, data in enumerate(tqdm.tqdm(train_loader, desc='Train')): total_steps += 1 model.set_input(data) model.optimize_parameters( check_grad=(opt.check_grad_freq > 0 and total_steps % opt.check_grad_freq == 0)) if total_steps % opt.display_freq == 0: train_error = model.get_current_errors() info = OrderedDict([ ('id', opt.id), ('iter', total_steps), ('epoch', epoch), ('lr', model.optimizers[0].param_groups[0]['lr']), ]) tqdm.tqdm.write(visualizer.log(info, train_error)) if epoch % opt.test_epoch_freq == 0: # model.get_current_errors() #erase training error information model.output = {} loss_buffer = LossBuffer(size=len(val_loader)) for i, data in enumerate(tqdm.tqdm(val_loader, desc='Test')): model.set_input(data) model.test(compute_loss=True) loss_buffer.add(model.get_current_errors()) test_error = loss_buffer.get_errors() info = OrderedDict([ ('time', time.ctime()), ('id', opt.id), ('epoch', epoch), ])
if opt.save_output: output_dir = os.path.join(model.save_dir, opt.output_dir) io.mkdir_if_missing(output_dir) total_time = 0 for i, data in enumerate(tqdm.tqdm(val_loader, desc='Test')): tic = time.time() model.eval() model.netG.eval() model.netF.eval() model.set_input(data) model.test() toc = time.time() total_time += (toc - tic) loss_buffer.add(model.get_current_errors()) # save output if opt.save_output: id_list = model.input['id'] images = model.output['img_out'].cpu().numpy().transpose(0, 2, 3, 1) images = ((images + 1.0) * 127.5).clip(0, 255).astype(np.uint8) for (sid1, sid2), img in zip(id_list, images): img = img[..., [2, 1, 0]] # convert to cv2 format cv2.imwrite(os.path.join(output_dir, '%s_%s.jpg' % (sid1, sid2)), img) test_error = loss_buffer.get_errors() test_error['sec_per_image'] = total_time / (opt.batch_size * len(val_loader)) info = OrderedDict([('model_id', opt.id), ('epoch', opt.which_epoch)]) log_str = visualizer.log(info, test_error, log_in_file=False) print(log_str)