if datapath.endswith('300W_LP'): pts = load_lua(osp.join(base_dir, f.split('_')[0], f[:-4] + '.t7'))[0] else: print("Given data set do not have 3D annotations.") exit() else: pts = load_lua(f) all_gts[i, :, :] = pts print('Loaded {} sample from {}'.format(len(lines), base_dir)) return all_gts, lines if __name__ == "__main__": import opts args = opts.argparser() dataset = args.data.split('/')[-1] save_dir = osp.join(args.checkpoint, dataset) print("save dictory: " + save_dir) preds = torch.from_numpy(loadpreds_if_exists(osp.join(save_dir, 'preds_valid.mat'))) gts, _ = loadgts(args.data, args.pointType) norm = np.ones(preds.size(0)) for i, gt in enumerate(gts): norm[i] = _get_bboxsize(gt) if dataset == 'LS3D-W' or dataset == '300VW-3D': for i in range(3): if dataset == 'LS3D-W': category = {'0': 'Easy', '1': 'Media', '2': 'Hard'}[str(i)] l, f = 2400*i, 2400*(i+1) else:
mean /= self.total std /= self.total ms = { 'mean': mean, 'std': std, } torch.save(ms, meanstd_file) if self.is_train: print('\tMean: %.4f, %.4f, %.4f' % (ms['mean'][0], ms['mean'][1], ms['mean'][2])) print('\tStd: %.4f, %.4f, %.4f' % (ms['std'][0], ms['std'][1], ms['std'][2])) return ms['mean'], ms['std'] if __name__ == "__main__": import opts, demo args = opts.argparser() dataset = W300LP(args, 'test') crop_win = None for i in range(dataset.__len__()): input, target, meta = dataset.__getitem__(i) input = input.numpy().transpose(1, 2, 0) * 255. target = target.numpy() if crop_win is None: crop_win = plt.imshow(input) else: crop_win.set_data(input) plt.pause(0.5) plt.draw
def main(): args = argparser() args.log_dir = os.path.join(args.log_path, 'log') args.checkpoint_dir = os.path.join(args.log_path, 'checkpoint') args.model_dir = os.path.join(args.log_path, 'model') check_log_folders(args) model = Model() # save model graph visualization model_writer = SummaryWriter(os.path.join(args.log_path, 'model')) inputs = torch.rand(8, 10, 2) model_writer.add_graph(model, inputs) model_writer.close() if args.device_ids != '' and torch.cuda.device_count() > 1: args.device_ids = args.device_ids.split(",") args.device_ids = list(map(lambda x: int(x), args.device_ids)) print("Let's use", torch.cuda.device_count(), "GPUs!") print(args.device_ids) torch.cuda.set_device(args.device_ids[0]) model = nn.DataParallel(model, device_ids=args.device_ids).cuda() else: args.device = torch.device("cuda:"+str(args.device) if torch.cuda.is_available() else "cpu") model.to(args.device) train_dataset = Dataset() val_dataset = Dataset() train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) optimizer = optim.SGD( model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) criterion = l2_loss scheduler = lr_scheduler.CosineAnnealingLR( optimizer, args.epochs, eta_min=0, last_epoch=-1 ) tf_writer = SummaryWriter(log_dir=os.path.join(args.log_path, 'log')) for epoch in range(args.epochs): # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args, tf_writer) # evaluate on validation set if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1: validate(val_loader, model, criterion, epoch*len(train_loader), args, tf_writer) save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, epoch, args) scheduler.step()