def get_csgnet():
csgnet = CsgNet(grid_shape=[64, 64, 64],
dropout=config.dropout,
mode=config.mode,
timesteps=7,
num_draws=len(unique_draws),
in_sz=config.input_size,
hd_sz=config.hidden_size,
stack_len=config.top_k)
weights = torch.load(config.pretrain_modelpath)
new_weights = {}
for k in weights.keys():
if k.startswith("module"):
new_weights[k[7:]] = weights[k]
csgnet.load_state_dict(new_weights)
csgnet.cuda()
for param in csgnet.parameters():
param.requires_grad = True
return csgnet
in_sz=config.input_size,
hd_sz=config.hidden_size,
stack_len=config.top_k)
# If you want to use multiple GPUs for training.
cuda_devices = torch.cuda.device_count()
if torch.cuda.device_count() > 1:
imitate_net.cuda_devices = torch.cuda.device_count()
print("using multi gpus", flush=True)
imitate_net = torch.nn.DataParallel(imitate_net, device_ids=[0, 1], dim=0)
imitate_net.cuda()
if config.preload_model:
imitate_net.load_state_dict(torch.load(config.pretrain_modelpath))
for param in imitate_net.parameters():
param.requires_grad = True
if config.optim == "sgd":
optimizer = optim.SGD(
[para for para in imitate_net.parameters() if para.requires_grad],
weight_decay=config.weight_decay,
momentum=0.9,
lr=config.lr,
nesterov=False)
elif config.optim == "adam":
optimizer = optim.Adam(
[para for para in imitate_net.parameters() if para.requires_grad],
weight_decay=config.weight_decay,
lr=config.lr)
print("using multi gpus", flush=True)
imitate_net = torch.nn.DataParallel(imitate_net, dim=0)
imitate_net.load_state_dict(torch.load(config.pretrain_modelpath))
else:
#weights = torch.load(config.pretrain_modelpath)
#new_weights = {}
#for k in weights.keys():
# if k.startswith("module"):
# new_weights[k[7:]] = weights[k]
#imitate_net.load_state_dict(new_weights)
imitate_net.load_state_dict(torch.load(config.pretrain_modelpath))
imitate_net.cuda()
for param in imitate_net.parameters():
param.requires_grad = True
config.test_size = sum(
(dataset_sizes[k][1] // config.batch_size) * config.batch_size
for k in dataset_sizes.keys())
for k in data_labels_paths.keys():
# if using multi gpu training, train and test batch size should be multiple of
# number of GPU edvices.
test_batch_size = config.batch_size
test_gen_objs[k] = generator.get_test_data(
test_batch_size,
k,
num_train_images=dataset_sizes[k][0],
num_test_images=dataset_sizes[k][1],