model.load_state_dict(torch.load(pretrained_weights))
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,
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)
