Exemplo n.º 1
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()

    # Create data loaders
    train_loader, test_loader = create_data_loaders(opt)

    # Create nn
    model, criterion_hm, criterion_paf, latest_inx = create_model(opt)
    model = model.cuda()
    criterion_hm = criterion_hm.cuda()
    criterion_paf = criterion_paf.cuda()

    # Create optimizer
    optimizer = create_optimizer(opt, model)

    # Other params
    n_epochs = opt.nEpoch
    to_train = opt.train
    drop_lr = opt.dropLR
    val_interval = opt.valInterval
    learn_rate = opt.LR
    visualize_out = opt.vizOut

    # train/ test
    train_net(train_loader, test_loader, model, criterion_hm, criterion_paf,
              optimizer, n_epochs, val_interval, learn_rate, drop_lr,
              opt.saveDir, visualize_out, latest_inx)
Exemplo n.º 2
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()

    # Create data loaders
    train_loader, test_loader = create_data_loaders(opt)

    # Create nn
    model, criterion_hm, criterion_paf, latest_inx = create_model(opt)
    # model = model.cuda()
    # criterion_hm = criterion_hm.cuda()
    # criterion_paf = criterion_paf.cuda()

    # Create optimizer
    optimizer = create_optimizer(opt, model)

    # Other params
    n_epochs = opt.nEpoch
    to_train = opt.train
    drop_lr = opt.dropLR
    val_interval = opt.valInterval
    learn_rate = opt.LR
    visualize_out = opt.vizOut

    # train/ test
    img, heat_map, paf, ignore_mask, keypoints = test_loader.dataset.get_item_raw(0, False)
Exemplo n.º 3
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()
    os.environ["CUDA_VISIBLE_DEVICES"] = opt["env"]["device"]
    print("Using GPU: {}".format(opt["env"]["device"]))
    # Create data loaders
    if opt["to_test"]:
        _, _, test_loader = create_data_loaders(opt)
    else:
        _, test_loader = create_data_loaders(opt)

    # Create nn
    model, _, _ = create_model(opt)
    model = model.cuda()
    # Get nn outputs
    outputs, indices = eval_net(test_loader, model, opt)

    if opt["val"]["dataset"] == 'coco':
        eval_COCO(outputs, opt["env"]["data"], indices, opt)
Exemplo n.º 4
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()

    # Create data loaders
    _, test_loader = create_data_loaders(opt)

    # Create nn
    model, _, _ = create_model(opt)
    model = model.cuda()

    # Get nn outputs
    outputs, indices = eval_net(test_loader, model, opt)

    if opt.dataset == 'coco':
        eval_COCO(outputs, opt.data, indices)
Exemplo n.º 5
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()
    opt.dataset = 'mine'
    opt.batchSize = 1
    opt.data = './data'
    opt.nThreads = 0
    opt.device = device
    fnames = ['test.jpg']

    # Create data loaders
    _, test_loader = create_data_loaders(opt, fnames)

    # Create nn
    model, _, _ = create_model(opt).to(opt.device)

    # output result images
    test_net(test_loader, model, opt)
Exemplo n.º 6
0
def main():
    # Seed all sources of randomness to 0 for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    random.seed(0)

    opt = Opts().parse()
    os.environ["CUDA_VISIBLE_DEVICES"] = opt["env"]["device"]
    print("Using GPU: {}".format(opt["env"]["device"]))

    # Create data loaders
    # Create data loaders
    train_loader, test_loader, _ = create_data_loaders(opt)
    # Create nn
    model, criterion_hm, criterion_paf = create_model(opt)
    model = torch.nn.DataParallel(model, device_ids=[int(index) for index in opt["env"]["device"].split(",")]).cuda() \
        if "," in opt["env"]["device"] else model.cuda()
    if opt["env"]["loadModel"] is not None and opt["typ"] == 'cpr':
        model.load_state_dict(torch.load(opt["env"]["loadModel"]))
        print('Loaded model from ' + opt["env"]["loadModel"])
    criterion_hm = criterion_hm.cuda()
    criterion_paf = criterion_paf.cuda()

    # Create optimizer
    optimizer = create_optimizer(opt, model)

    # Other params
    to_train = opt["to_train"]
    visualize_out = opt["viz"]["vizOut"]

    # train/ test
    Processer = process(model)
    if to_train:
        Processer.train_net(train_loader, test_loader, criterion_hm, criterion_paf, optimizer, opt, viz_output=visualize_out)
    else:
        Processer.validate_net(test_loader, criterion_hm, criterion_paf, save_dir=opt["env"]["saveDir"], viz_output=visualize_out)