# generator
model = FullNetwork(vp_value_count=VP_VALUE_COUNT, stdev=STDEV,
output_shape=(BATCH_SIZE, CHANNELS, FRAMES, HEIGHT, WIDTH), use_est_vp=False)
model = model.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LR)
# data
trainset = PanopticDataset(root_dir=data_root_dir, data_file=train_split,
resize_height=HEIGHT, resize_width=WIDTH,
clip_len=FRAMES, skip_len=SKIP_LEN,
random_all=RANDOM_ALL, close_views=CLOSE_VIEWS,
close_cams_file=close_cams_file, precrop=PRECROP)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
testset = PanopticDataset(root_dir=data_root_dir, data_file=test_split,
resize_height=HEIGHT, resize_width=WIDTH,
clip_len=FRAMES, skip_len=SKIP_LEN,
random_all=RANDOM_ALL, close_views=CLOSE_VIEWS,
close_cams_file=close_cams_file, precrop=PRECROP)
testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
else:
print('This network has only been set up to train on the NTU and panoptic datasets.')
print_params()
model = FullNetwork(vp_value_count=3,
output_shape=(BATCH_SIZE, CHANNELS, FRAMES, HEIGHT,
WIDTH))
model.load_state_dict(torch.load(weights_path))
model = model.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
criterion = nn.MSELoss()
testset = PanopticDataset(root_dir=data_root_dir,
data_file=test_split,
resize_height=HEIGHT,
resize_width=WIDTH,
clip_len=FRAMES,
skip_len=SKIP_LEN,
random_all=RANDOM_ALL,
precrop=PRECROP)
testloader = torch.utils.data.DataLoader(testset,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=2)
else:
print(
'This network has only been set up to run on the NTU and panoptic datasets.'
)
print_params()
print(model)