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
7: [proportion * 4000, proportion * 100]
}
config.train_size = sum(dataset_sizes[k][0] for k in dataset_sizes.keys())
config.test_size = sum(dataset_sizes[k][1] for k in dataset_sizes.keys())
types_prog = len(dataset_sizes)
generator = Generator(data_labels_paths=data_labels_paths,
batch_size=config.batch_size,
time_steps=max(data_labels_paths.keys()),
stack_size=max(data_labels_paths.keys()) // 2 + 1)
imitate_net = CsgNet(grid_shape=[64, 64, 64],
dropout=config.dropout,
mode=config.mode,
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))
config.train_size = sum(dataset_sizes[k][0] for k in dataset_sizes.keys())
config.test_size = sum(dataset_sizes[k][1] for k in dataset_sizes.keys())
types_prog = len(dataset_sizes)
# primitives = dd.io.load('data/primitives_cuboids.h5')
primitives = None
generator = Generator(data_labels_paths=data_labels_paths,
batch_size=config.batch_size,
time_steps=max(data_labels_paths.keys()),
stack_size=max(data_labels_paths.keys()) // 2 + 1,
primitives=primitives)
imitate_net = CsgNet(grid_shape=[64, 64, 64],
dropout=config.dropout,
mode=config.mode,
timesteps=max(data_labels_paths.keys()),
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))
}
dataset_sizes = {3: [110000, 20000], 5: [220000, 40000], 7: [440000, 80000]}
test_gen_objs = {}
types_prog = len(dataset_sizes.keys())
max_len = max(data_labels_paths.keys())
generator = Generator(data_labels_paths=data_labels_paths,
batch_size=config.batch_size,
time_steps=max(data_labels_paths.keys()),
stack_size=max(data_labels_paths.keys()) // 2 + 1)
imitate_net = CsgNet(grid_shape=[64, 64, 64],
dropout=config.dropout,
mode=config.mode,
timesteps=max_len,
num_draws=len(generator.unique_draw),
in_sz=config.input_size,
hd_sz=config.hidden_size,
stack_len=config.top_k)
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, 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"):
7: [440000, 80000]}
test_gen_objs = {}
types_prog = len(dataset_sizes.keys())
max_len = max(data_labels_paths.keys())
beam_width = 10
stack_size = max_len // 2 + 1
generator = Generator(data_labels_paths=data_labels_paths,
batch_size=config.batch_size,
time_steps=max(data_labels_paths.keys()),
stack_size=max(data_labels_paths.keys()) // 2 + 1)
imitate_net = CsgNet(grid_shape=[64, 64, 64], dropout=config.dropout,
mode=config.mode, timesteps=max_len,
num_draws=len(generator.unique_draw),
in_sz=config.input_size,
hd_sz=config.hidden_size,
stack_len=config.top_k)
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, 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"):