def main(argv): (opts, args) = parser.parse_args(argv) # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] max_iterations = config.hyperparameters['max_iterations'] train_loader_a = get_data_loader(config.datasets['train_a'], batch_size) train_loader_b = get_data_loader(config.datasets['train_b'], batch_size) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters['trainer'] local_dict = locals() exec(cmd,globals(),local_dict) trainer = local_dict['trainer'] # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix) trainer.cuda(opts.gpu) ###################################################################################################################### # Setup logger and repare image outputs train_writer = tensorboard.FileWriter("%s/%s" % (opts.log,os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder(config.snapshot_prefix, iterations, config.image_save_iterations) for ep in range(0, MAX_EPOCHS): for it, (images_a, images_b) in enumerate(izip(train_loader_a,train_loader_b)): if images_a.size(0) != batch_size or images_b.size(0) != batch_size: continue images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) # Main training code trainer.dis_update(images_a, images_b, config.hyperparameters) image_outputs = trainer.gen_update(images_a, images_b, config.hyperparameters) assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs) # Dump training stats in log file if (iterations+1) % config.display == 0: write_loss(iterations, max_iterations, trainer, train_writer) if (iterations+1) % config.image_save_iterations == 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) write_html(snapshot_directory + "/index.html", iterations + 1, config.image_save_iterations, image_directory) elif (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) # Save network weights if (iterations+1) % config.snapshot_save_iterations == 0: trainer.save(config.snapshot_prefix, iterations) iterations += 1 if iterations >= max_iterations: return
def dcgan_train(config, checkpoint_dir=None): step = 0 use_cuda = config.get("use_gpu") and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") netD = Discriminator().to(device) netD.apply(weights_init) netG = Generator().to(device) netG.apply(weights_init) criterion = nn.BCELoss() optimizerD = optim.Adam(netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) with FileLock(os.path.expanduser("~/.data.lock")): dataloader = get_data_loader() if checkpoint_dir is not None: path = os.path.join(checkpoint_dir, "checkpoint") checkpoint = torch.load(path) netD.load_state_dict(checkpoint["netDmodel"]) netG.load_state_dict(checkpoint["netGmodel"]) optimizerD.load_state_dict(checkpoint["optimD"]) optimizerG.load_state_dict(checkpoint["optimG"]) step = checkpoint["step"] if "netD_lr" in config: for param_group in optimizerD.param_groups: param_group["lr"] = config["netD_lr"] if "netG_lr" in config: for param_group in optimizerG.param_groups: param_group["lr"] = config["netG_lr"] while True: lossG, lossD, is_score = train( netD, netG, optimizerG, optimizerD, criterion, dataloader, step, device, config["mnist_model_ref"], ) step += 1 with tune.checkpoint_dir(step=step) as checkpoint_dir: path = os.path.join(checkpoint_dir, "checkpoint") torch.save( { "netDmodel": netD.state_dict(), "netGmodel": netG.state_dict(), "optimD": optimizerD.state_dict(), "optimG": optimizerG.state_dict(), "step": step, }, path, ) tune.report(lossg=lossG, lossd=lossD, is_score=is_score)
def setup(self, config): use_cuda = config.get("use_gpu") and torch.cuda.is_available() self.device = torch.device("cuda" if use_cuda else "cpu") self.netD = Discriminator().to(self.device) self.netD.apply(weights_init) self.netG = Generator().to(self.device) self.netG.apply(weights_init) self.criterion = nn.BCELoss() self.optimizerD = optim.Adam(self.netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) self.optimizerG = optim.Adam(self.netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) with FileLock(os.path.expanduser("~/.data.lock")): self.dataloader = get_data_loader(config.get("data_dir", "~/data")) self.mnist_model_ref = config["mnist_model_ref"]
parser.add_argument('--glove', default='glove/glove.6B.100d.txt', help='path to glove txt') parser.add_argument('--model', default='rnn', help='model') parser.add_argument('--rnn', default='LSTM', choices=['LSTM', 'GRU'], help='rnn module type') parser.add_argument('--mean_seq', default=False, action='store_true', help='use mean of rnn output') parser.add_argument('--clip', type=float, default=0.25, help='gradient clipping') parser.add_argument('--max-sentence-length', type=int, default=256, help='default length of arg1:arg2') parser.add_argument('--pdtb-category', default='Comparison', choices=['Comparison', 'Contingency', 'Temporal', 'Expansion', ''], help='PDTB category') args = parser.parse_args() d_word_index, results_path = common.get_word_index(args.model, args.glove, args.embedding_size) # create tester print("===> creating dataloaders ...") val_loader = common.get_data_loader(args.model, 'test', d_word_index, args.batch_size, args.max_sentence_length, pdtb_category=args.pdtb_category) # load model,optimizer and loss model, optimizer, criterion = common.get_model(model=args.model, model_path=results_path, lr=args.lr, weight_decay=args.weight_decay, pdtb_category=args.pdtb_category) print(optimizer) print(criterion) if args.cuda: torch.backends.cudnn.enabled = True cudnn.benchmark = True model.cuda() criterion = criterion.cuda()
parser.add_argument( "--smoke-test", action="store_true", help="Finish quickly for testing") args, _ = parser.parse_known_args() ray.init() import urllib.request # Download a pre-trained MNIST model for inception score calculation. # This is a tiny model (<100kb). if not os.path.exists(MODEL_PATH): print("downloading model") os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True) urllib.request.urlretrieve( "https://github.com/ray-project/ray/raw/master/python/ray/tune/" "examples/pbt_dcgan_mnist/mnist_cnn.pt", MODEL_PATH) dataloader = get_data_loader() if not args.smoke_test: plot_images(dataloader) # load the pretrained mnist classification model for inception_score mnist_cnn = Net() mnist_cnn.load_state_dict(torch.load(MODEL_PATH)) mnist_cnn.eval() mnist_model_ref = ray.put(mnist_cnn) # __tune_begin__ scheduler = PopulationBasedTraining( time_attr="training_iteration", metric="is_score", mode="max", perturbation_interval=5,
print(ckpt_path) for path in [base_path, exp_path, data_path, base_path, ckpt_path, log_path]: if not os.path.exists(path): os.makedirs(path) use_cuda = torch.cuda.is_available() if use_cuda: torch.cuda.set_device(args.gpu_id) best_acc = 0 # best test accuracy start_epoch = 0 # start from epoch 0 or last checkpoint epoch # Data print('==> Preparing data..') config = NetConfig(args.config) trainloader = get_data_loader(config.datasets['traindata'], args.batch_size) testloader = get_data_loader(config.datasets['valdata'], args.batch_size) # Model if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') #assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!' checkpoint = torch.load('%s/ckpt.t7' % ckpt_path) net = checkpoint['net'] best_acc = checkpoint['acc'] start_epoch = checkpoint['epoch'] else: net = VGG('VGG11', 4) #from torchvision.models.vgg import * #net = vgg11_bn(num_classes=4) tensorboard_logger.configure(log_path)
def dcgan_train(config): step = 0 use_cuda = config.get("use_gpu") and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") netD = Discriminator().to(device) netD.apply(weights_init) netG = Generator().to(device) netG.apply(weights_init) criterion = nn.BCELoss() optimizerD = optim.Adam(netD.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=config.get("lr", 0.01), betas=(beta1, 0.999)) with FileLock(os.path.expanduser("~/.data.lock")): dataloader = get_data_loader() if session.get_checkpoint(): loaded_checkpoint = session.get_checkpoint() with loaded_checkpoint.as_directory() as loaded_checkpoint_dir: path = os.path.join(loaded_checkpoint_dir, "checkpoint.pt") checkpoint = torch.load(path) netD.load_state_dict(checkpoint["netDmodel"]) netG.load_state_dict(checkpoint["netGmodel"]) optimizerD.load_state_dict(checkpoint["optimD"]) optimizerG.load_state_dict(checkpoint["optimG"]) step = checkpoint["step"] if "netD_lr" in config: for param_group in optimizerD.param_groups: param_group["lr"] = config["netD_lr"] if "netG_lr" in config: for param_group in optimizerG.param_groups: param_group["lr"] = config["netG_lr"] while True: lossG, lossD, is_score = train( netD, netG, optimizerG, optimizerD, criterion, dataloader, step, device, config["mnist_model_ref"], ) step += 1 os.makedirs("my_model", exist_ok=True) torch.save( { "netDmodel": netD.state_dict(), "netGmodel": netG.state_dict(), "optimD": optimizerD.state_dict(), "optimG": optimizerG.state_dict(), "step": step, }, "my_model/checkpoint.pt", ) session.report( { "lossg": lossG, "lossd": lossD, "is_score": is_score }, checkpoint=Checkpoint.from_directory("my_model"), )
ray.init() import urllib.request # Download a pre-trained MNIST model for inception score calculation. # This is a tiny model (<100kb). if not os.path.exists(MODEL_PATH): print("downloading model") os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True) urllib.request.urlretrieve( "https://github.com/ray-project/ray/raw/master/python/ray/tune/" "examples/pbt_dcgan_mnist/mnist_cnn.pt", MODEL_PATH, ) dataloader = get_data_loader(args.data_dir) if not args.smoke_test: plot_images(dataloader) # __tune_begin__ # load the pretrained mnist classification model for inception_score mnist_cnn = Net() mnist_cnn.load_state_dict(torch.load(MODEL_PATH)) mnist_cnn.eval() # Put the model in Ray object store. mnist_model_ref = ray.put(mnist_cnn) scheduler = PopulationBasedTraining( perturbation_interval=5, hyperparam_mutations={
def main(argv): (opts, args) = parser.parse_args(argv) if 'estimate' in opts.mode: mode_idx = int(opts.mode[-1]) global colorPlatte, bones, Evaluation if 'nyu' in opts.config: colorPlatte = utils.util.nyuColorIdx bones = utils.util.nyuBones Evaluation = NYUHandposeEvaluation elif 'icvl' in opts.config: colorPlatte = utils.util.icvlColorIdx bones = utils.util.icvlBones Evaluation = ICVLHandposeEvaluation # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] if 'estimate' in opts.mode else 1 test_batch_size = batch_size * 32 max_iterations = config.hyperparameters['max_iterations'] frac = opts.frac dataset_a = get_dataset(config.datasets['train_a']) dataset_b = get_dataset(config.datasets['train_b']) dataset_test = get_dataset(config.datasets['test_b']) train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True) train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True) test_loader_real = get_data_loader(dataset_test, test_batch_size, shuffle=False) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters['trainer'] local_dict = locals() exec(cmd,globals(),local_dict) trainer = local_dict['trainer'] di_a = dataset_a.di di_b = dataset_b.di # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix, idx=-1, load_opt=True) for i in range(iterations//1000): trainer.dis_sch.step() trainer.gen_sch.step() trainer.cuda(opts.gpu) print('using %.2f percent of the labeled real data' % frac) try: if 'estimate' in opts.mode and (mode_idx == 3 or mode_idx == 4): trainer.load_vae(config.snapshot_prefix, 2+frac) else: trainer.load_vae(config.snapshot_prefix, frac) except: print('Failed to load the parameters of vae') if 'estimate' in opts.mode: if opts.idx != 0: trainer.resume(config.snapshot_prefix, idx=opts.idx, est=mode_idx==5) if frac > 0. and frac < 1.: dataset_b.set_nmax(frac) #trainer.dis.freeze_layers() ############################################################################################### # Setup logger and repare image outputs train_writer = tensorboardX.FileWriter("%s/%s" % (opts.log,os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder(config.snapshot_prefix, iterations, config.image_save_iterations) best_err, best_acc = 100., 0. start_time = time.time() for ep in range(0, MAX_EPOCHS): for it, ((images_a, labels_a, com_a, M_a, cube_a, _), (images_b,labels_b, com_b, M_b, cube_b, _)) in \ enumerate(izip(train_loader_a,train_loader_b)): if images_a.size(0) != batch_size or images_b.size(0) != batch_size: continue images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) labels_a = Variable(labels_a.cuda(opts.gpu)) labels_b = Variable(labels_b.cuda(opts.gpu)) com_a = Variable(com_a.cuda(opts.gpu)) com_b = Variable(com_b.cuda(opts.gpu)) trainer.dis.train() if opts.mode == 'pretrain': if (iterations+1) % 1000 == 0: trainer.dis_sch.step() trainer.gen_sch.step() print('lr %.8f' % trainer.dis_sch.get_lr()[0]) trainer.dis_update(images_a, labels_a, images_b, labels_b, com_a, com_b, config.hyperparameters) image_outputs = trainer.gen_update(images_a, labels_a, images_b, labels_b, config.hyperparameters) assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs) else: if (iterations+1) % 100 == 0: trainer.dis_sch.step() image_outputs = trainer.post_update(images_a, labels_a, images_b, labels_b,com_a,com_b, mode_idx, config.hyperparameters) assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs) # Dump training stats in log file if (iterations+1) % config.display == 0: elapsed_time = time.time() - start_time write_loss(iterations, max_iterations, trainer, train_writer, elapsed_time) start_time = time.time() if (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) if (iterations+1) % config.image_save_iterations == 0: if opts.mode == 'pretrain':# and (iterations+1) % (2*config.image_save_iterations) != 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) write_html(snapshot_directory + "/index.html", iterations + 1, \ config.image_save_iterations, image_directory) else: trainer.dis.eval() score, maxerr = 0, 0 num_samples = 0 maxJntError = [] meanJntError = 0 img2sav = None gt3D = [] joints = [] joints_imgcord = [] codec = cv2.VideoWriter_fourcc(*'XVID') vid = cv2.VideoWriter(os.path.join(image_directory,'gen.avi'), codec, 25, (128*2,128)) for tit, (test_images_b, test_labels_b, com_b, trans_b, cube_b, fn) in enumerate(test_loader_real): test_images_b = Variable(test_images_b.cuda(opts.gpu)) test_labels_b = Variable(test_labels_b.cuda(opts.gpu)) if mode_idx == 0: pred_pose, pred_post, _ = trainer.dis.regress_a(test_images_b) else: pred_pose, pred_post, _ = trainer.dis.regress_b(test_images_b) if True: pred_pose = trainer.vae.decode(pred_post) n = test_labels_b.size(0) gt_pose = test_labels_b.data.cpu().numpy().reshape((n,-1, 3)) pr_pose = pred_pose.data.cpu().numpy().reshape((n,-1, 3)) if tit < 20: for i in range(0, n, 4): real_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), gt_pose[i].reshape((-1)), \ trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0) est_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), pr_pose[i].reshape((-1)), \ trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0) vid.write(np.hstack((real_img,est_img)).astype('uint8')) both_img = np.vstack((real_img,est_img)) if True and tit < 8: if img2sav is None: img2sav = both_img else: img2sav = np.hstack((img2sav,both_img)) if 'nyu' in opts.config: restrictedJointsEval = np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 27, 30, 31, 32]) gt_pose = gt_pose[:,restrictedJointsEval] pr_pose = pr_pose[:,restrictedJointsEval] for i in range(n): gt3D.append(gt_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy()) joints.append(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy()) joints_imgcord.append(di_b.joints3DToImg(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy())) score += meanJntError num_samples += test_images_b.size(0) cv2.imwrite(image_directory + '/_test.jpg', img2sav.astype('uint8')) vid.release() hpe = Evaluation(np.array(gt3D), np.array(joints)) mean_err = hpe.getMeanError() over_40 = 100. * hpe.getNumFramesWithinMaxDist(40) / len(gt3D) best_err = np.minimum(best_err, mean_err) best_acc = np.maximum(best_acc, over_40) print("------------ Mean err: {:.4f} ({:.4f}) mm, Max over 40mm: {:.2f} ({:.2f}) %".format(mean_err, best_err, over_40, best_acc)) # Save network weights if (iterations+1) % config.snapshot_save_iterations == 0: if opts.mode == 'pretrain': trainer.save(config.snapshot_prefix, iterations) elif 'estimate' in opts.mode: trainer.save(config.snapshot_prefix+'_est', iterations) iterations += 1 if iterations >= max_iterations: return
formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--conf', default='conf/default.json', help='path to conf file') parsed_args = parser.parse_args() args = json.load(open(parsed_args.conf)) d_word_index, results_path = common.get_word_index(args['model'], args['glove'], args['embedding_size']) # create tester print("===> creating dataloaders ...") val_loader = common.get_data_loader(args['model'], 'test', d_word_index, args['batch_size'], args['max_sentence_length'], pdtb_category=args['pdtb_category']) # load model,optimizer and loss model, optimizer, criterion = common.get_model( model=args['model'], model_path=results_path, lr=args['learning_rate'], weight_decay=args['weight_decay'], pdtb_category=args['pdtb_category']) print(optimizer) print(criterion) if args['cuda']: torch.backends.cudnn.enabled = True
def main(argv): (opts, args) = parser.parse_args(argv) # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] max_iterations = config.hyperparameters['max_iterations'] train_loader_a = get_data_loader(config.datasets['train_a'], batch_size) train_loader_b = get_data_loader(config.datasets['train_b'], batch_size) # Parse ROI parameters roi = [int(val_str) for val_str in opts.roi.split(',')] roi_x = roi[0] roi_y = roi[1] roi_w = roi[2] roi_h = roi[3] cmd1 = "trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] cmd2 = "roi_trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] local_dict = locals() exec(cmd1, globals(), local_dict) trainer = local_dict['trainer'] exec(cmd2, globals(), local_dict) roi_trainer = local_dict['roi_trainer'] # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix) roi_trainer.resume(config.snapshot_prefix) trainer.cuda(opts.gpu) roi_trainer.cuda(opts.gpu) ###################################################################################################################### # Setup logger and repare image outputs train_writer = tensorboard.FileWriter( "%s/%s" % (opts.log, os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder( config.snapshot_prefix, iterations, config.image_save_iterations) for ep in range(0, MAX_EPOCHS): for it, (images_a, images_b) in enumerate(izip(train_loader_a, train_loader_b)): if images_a.size(0) != batch_size or images_b.size( 0) != batch_size: continue # Crop images according to ROI roi_images_a = images_a[:, :, roi_y:roi_y + roi_h, roi_x:roi_x + roi_w].clone() roi_images_b = images_b[:, :, roi_y:roi_y + roi_h, roi_x:roi_x + roi_w].clone() roi_images_a = Variable(roi_images_a.cuda(opts.gpu)) roi_images_b = Variable(roi_images_b.cuda(opts.gpu)) images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) # Main training code trainer.dis_update(images_a, images_b, config.hyperparameters) trainer.gen_update(images_a, images_b, config.hyperparameters) # Training code for ROI roi_trainer.dis_update(roi_images_a, roi_images_b, config.hyperparameters) roi_image_outputs = roi_trainer.gen_update(roi_images_a, roi_images_b, config.hyperparameters) roi_assembled_images = roi_trainer.assemble_outputs( roi_images_a, roi_images_b, roi_image_outputs) # Paste ROI to original images to update generator x_aa, x_ba, x_ab, x_bb, shared = trainer.gen(images_a, images_b) x_ba_paste = x_ba.clone() x_ab_paste = x_ab.clone() x_ba_paste[:, :, roi_y:roi_y + roi_h, roi_x:roi_x + roi_w] = roi_image_outputs[1].clone() x_ab_paste[:, :, roi_y:roi_y + roi_h, roi_x:roi_x + roi_w] = roi_image_outputs[2].clone() trainer.gen.zero_grad() image_outputs = trainer.gen_update_helper(images_a, images_b, x_aa, x_ba_paste, x_ab_paste, x_bb, shared, config.hyperparameters) assembled_images = trainer.assemble_outputs( images_a, images_b, image_outputs) # Dump training stats in log file if (iterations + 1) % config.display == 0: write_loss(iterations, max_iterations, trainer, train_writer) if (iterations + 1) % config.image_save_iterations == 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) img_filename = '%s/roi_gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(roi_assembled_images.data / 2 + 0.5, img_filename, nrow=1) write_html(snapshot_directory + "/index.html", iterations + 1, config.image_save_iterations, image_directory) elif (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) img_filename = '%s/roi_gen.jpg' % (image_directory) torchvision.utils.save_image(roi_assembled_images.data / 2 + 0.5, img_filename, nrow=1) # Save network weights if (iterations + 1) % config.snapshot_save_iterations == 0: trainer.save(config.snapshot_prefix, iterations) iterations += 1 if iterations >= max_iterations: return
def main(argv): (opts, args) = parser.parse_args(argv) seed = 0 torch.cuda.manual_seed(seed) # Set seed for deterministic, is this enough? torch.manual_seed(seed) np.random.seed(seed=seed) # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] max_iterations = config.hyperparameters['max_iterations'] train_loader_a = get_data_loader(config.datasets['train_a'], batch_size) train_loader_b = get_data_loader(config.datasets['train_b'], batch_size) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] local_dict = locals() exec(cmd, globals(), local_dict) trainer = local_dict['trainer'] # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix) trainer.cuda(opts.gpu) ###################################################################################################################### # Setup logger and repare image outputs train_writer = tf.summary.FileWriter( "%s/%s" % (opts.log, os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder( config.snapshot_prefix, iterations, config.image_save_iterations) for ep in range(0, MAX_EPOCHS): for it, (data_a, data_b) in enumerate(izip(train_loader_a, train_loader_b)): images_a = data_a['data'] labels_a = data_a.get("data_lab") images_b = data_b['data'] labels_b = data_b.get("data_lab") if images_a.size(0) != batch_size or images_b.size( 0) != batch_size: continue images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) # Main training code trainer.dis_update(images_a, images_b, config.hyperparameters) image_outputs = trainer.gen_update(images_a, images_b, config.hyperparameters, labels_a, labels_b) assembled_images = trainer.assemble_outputs( images_a, images_b, image_outputs) # Dump training stats in log file if (iterations + 1) % config.display == 0: write_loss(iterations, max_iterations, trainer, train_writer) if (iterations + 1) % config.image_save_iterations == 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data, img_filename, nrow=1) segm_image = None if labels_a is not None: _, enet_classes_ab = torch.max(image_outputs[7], dim=1, keepdim=False) segm_image_ab = np.concatenate( (np.squeeze(enet_classes_ab.data.cpu().numpy()), np.squeeze(labels_a.cpu().numpy())), axis=1) segm_image = segm_image_ab if labels_b is not None: _, enet_classes_ba = torch.max(image_outputs[6], dim=1, keepdim=False) segm_image_ba = np.concatenate( (np.squeeze(enet_classes_ba.data.cpu().numpy()), np.squeeze(labels_b.cpu().numpy())), axis=1) if labels_a is not None: segm_image = np.concatenate( (segm_image_ba, segm_image_ab), axis=0) else: segm_image = segm_image_ba if segm_image is not None: segm_filename = '%s/segm_cat_%08d.jpg' % (image_directory, iterations + 1) cv2.imwrite(segm_filename, segm_image) write_html(snapshot_directory + "/index.html", iterations + 1, config.image_save_iterations, image_directory) elif (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data, img_filename, nrow=1) # Save network weights if (iterations + 1) % config.snapshot_save_iterations == 0: trainer.save(config.snapshot_prefix, iterations) iterations += 1 if iterations >= max_iterations: return del images_a del images_b del image_outputs del assembled_images del labels_a del labels_b
def main(argv): (opts, args) = parser.parse_args(argv) # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] max_iterations = config.hyperparameters['max_iterations'] train_loader_a = get_data_loader(config.datasets['train_a'], batch_size) train_loader_b = get_data_loader(config.datasets['train_b'], batch_size) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] local_dict = locals() exec(cmd, globals(), local_dict) trainer = local_dict['trainer'] # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix) trainer.cuda(opts.gpu) #trainer = torch.nn.DataParallel(trainer, device_ids=range(torch.cuda.device_count())) # replicates the model in all the GPUs. The batch size should be a multiple of the number of GPUs. This doesn't work since it complains about dis_uptate is not member of DataParallel. cudnn.benchmark = True ###################################################################################################################### # Setup logger and repare image outputs train_writer = SummaryWriter( "%s/%s" % (opts.log, os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder( config.snapshot_prefix, iterations, config.image_save_iterations) for ep in range(0, MAX_EPOCHS): for it, (images_a, images_b) in enumerate(izip(train_loader_a, train_loader_b)): if images_a.size(0) != batch_size or images_b.size( 0) != batch_size: continue images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) # Main training code trainer.dis_update(images_a, images_b, config.hyperparameters) image_outputs = trainer.gen_update(images_a, images_b, config.hyperparameters) assembled_images = trainer.assemble_outputs( images_a, images_b, image_outputs) # Dump training stats in log file if (iterations + 1) % config.display == 0: write_loss_X(iterations, max_iterations, trainer, train_writer) if (iterations + 1) % config.image_save_iterations == 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) write_html(snapshot_directory + "/index.html", iterations + 1, config.image_save_iterations, image_directory) elif (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) # Save network weights if (iterations + 1) % config.snapshot_save_iterations == 0: trainer.save(config.snapshot_prefix, iterations) iterations += 1 if iterations >= max_iterations: return
def main(argv): (opts, args) = parser.parse_args(argv) global colorPlatte, bones, Evaluation if 'nyu' in opts.config: colorPlatte = utils.util.nyuColorIdx bones = utils.util.nyuBones Evaluation = NYUHandposeEvaluation elif 'icvl' in opts.config: colorPlatte = utils.util.icvlColorIdx bones = utils.util.icvlBones Evaluation = ICVLHandposeEvaluation # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size_pose'] max_iterations = 200000 #config.hyperparameters['max_iterations'] frac = opts.frac dataset_a = get_dataset(config.datasets['train_a']) dataset_b = get_dataset(config.datasets['train_b']) dataset_test = get_dataset(config.datasets['test_b']) train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True) train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True) test_loader_real = get_data_loader(dataset_test, 1, shuffle=True) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] local_dict = locals() exec(cmd, globals(), local_dict) trainer = local_dict['trainer'] iterations = 0 trainer.cuda(opts.gpu) dataset_a.pose_only = True dataset_b.pose_only = True if frac > 0. and frac < 1.: dataset_b.set_nmax(frac) di_a = dataset_a.di di_b = dataset_b.di dataset_a.sample_poses() dataset_b.sample_poses() ################################################################### # Setup logger and repare image outputs train_writer = tensorboardX.FileWriter( "%s/%s" % (opts.log, os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder( config.snapshot_prefix, iterations, config.image_save_iterations) print('using %.2f percent of the labeled real data' % frac) start_time = time.time() for ep in range(0, MAX_EPOCHS): for it, ((labels_a), (labels_b)) in enumerate(izip(train_loader_a, train_loader_b)): if labels_a.size(0) != batch_size or labels_b.size( 0) != batch_size: continue labels_a = Variable(labels_a.cuda(opts.gpu)) labels_b = Variable(labels_b.cuda(opts.gpu)) labels = labels_a if frac > 0.: labels = torch.cat((labels_a, labels_b), 0) if (iterations + 1) % 1000 == 0: trainer.vae_sch.step() recon_pose = trainer.vae_update(labels, config.hyperparameters) # Dump training stats in log file if (iterations + 1) % config.display == 0: elapsed_time = time.time() - start_time write_loss(iterations, max_iterations, trainer, train_writer, elapsed_time) start_time = time.time() if (iterations + 1) % (10 * config.image_save_iterations) == 0: if True: score, maxerr = 0, 0 num_samples = 0 maxJntError = [] img2sav = None gt3D = [] joints = [] for tit, (test_images_b, test_labels_b, com_b, trans_b, cube_b, _) in enumerate(test_loader_real): test_images_b = Variable(test_images_b.cuda(opts.gpu)) test_labels_b = Variable(test_labels_b.cuda(opts.gpu)) pred_pose = trainer.vae.decode( trainer.vae.encode(test_labels_b)[1]) gt3D.append(test_labels_b.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\ com_b.numpy()) joints.append(pred_pose.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\ com_b.numpy()) if True and tit < 8: real_img = visPair(di_b, test_images_b.data.cpu().numpy(), test_labels_b.data.cpu().numpy(), \ trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0) est_img = visPair(di_b, test_images_b.data.cpu().numpy(), pred_pose.data.cpu().numpy(), \ trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0) if img2sav is None: img2sav = np.vstack((real_img, est_img)) else: img2sav = np.hstack( (img2sav, np.vstack((real_img, est_img)))) num_samples += test_images_b.size(0) cv2.imwrite(image_directory + '/_test.jpg', img2sav.astype('uint8')) #maxerr = Evaluation.plotError(maxJntError, image_directory + '/maxJntError.txt') hpe = Evaluation(np.array(gt3D), np.array(joints)) print("Mean error: {}mm, max error: {}mm".format( hpe.getMeanError(), hpe.getMaxError())) # Save network weights if (iterations + 1) % (4 * config.snapshot_save_iterations) == 0: trainer.save_vae(config.snapshot_prefix, iterations, 2 + frac) iterations += 1 if iterations >= max_iterations: return