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
timesteps=max(data_labels_paths.keys()),
num_draws=len(generator.unique_draw),
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],