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"):