def get_csgnet(): config = read_config.Config("config_synthetic.yml") # Encoder encoder_net = Encoder(config.encoder_drop) encoder_net = encoder_net.to(device) imitate_net = ImitateJoint(hd_sz=config.hidden_size, input_size=config.input_size, encoder=encoder_net, mode=config.mode, num_draws=400, canvas_shape=config.canvas_shape) imitate_net = imitate_net.to(device) print("pre loading model") pretrained_dict = torch.load(config.pretrain_modelpath, map_location=device) imitate_net_dict = imitate_net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in imitate_net_dict } imitate_net_dict.update(pretrained_dict) imitate_net.load_state_dict(imitate_net_dict) return imitate_net.encoder
def get_blank_csgnet(): config = read_config.Config("config_synthetic.yml") # Encoder encoder_net = Encoder(config.encoder_drop) encoder_net = encoder_net.to(device) imitate_net = ImitateJoint(hd_sz=config.hidden_size, input_size=config.input_size, encoder=encoder_net, mode=config.mode, num_draws=400, canvas_shape=config.canvas_shape) imitate_net = imitate_net.to(device) return imitate_net
def get_csgnet(): config = read_config.Config("config_synthetic.yml") # Encoder encoder_net = Encoder(config.encoder_drop) encoder_net = encoder_net.to(device) # Load the terminals symbols of the grammar with open("terminals.txt", "r") as file: unique_draw = file.readlines() for index, e in enumerate(unique_draw): unique_draw[index] = e[0:-1] imitate_net = ImitateJoint(hd_sz=config.hidden_size, input_size=config.input_size, encoder=encoder_net, mode=config.mode, num_draws=len(unique_draw), canvas_shape=config.canvas_shape) imitate_net = imitate_net.to(device) print("pre loading model") pretrained_dict = torch.load(config.pretrain_modelpath, map_location=device) imitate_net_dict = imitate_net.state_dict() imitate_pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in imitate_net_dict } imitate_net_dict.update(imitate_pretrained_dict) imitate_net.load_state_dict(imitate_net_dict) for param in imitate_net.parameters(): param.requires_grad = True for param in encoder_net.parameters(): param.requires_grad = True return (encoder_net, imitate_net)
logger.info(config.config) # CNN encoder encoder_net = Encoder(config.encoder_drop) encoder_net.cuda() # Load the terminals symbols of the grammar with open("terminals.txt", "r") as file: unique_draw = file.readlines() for index, e in enumerate(unique_draw): unique_draw[index] = e[0:-1] # RNN decoder imitate_net = ImitateJoint(hd_sz=config.hidden_size, input_size=config.input_size, encoder=encoder_net, mode=config.mode, num_draws=len(unique_draw), canvas_shape=config.canvas_shape) imitate_net.cuda() imitate_net.epsilon = config.eps if config.preload_model: print("pre loading model") pretrained_dict = torch.load(config.pretrain_modelpath) imitate_net_dict = imitate_net.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in imitate_net_dict } imitate_net_dict.update(pretrained_dict) imitate_net.load_state_dict(imitate_net_dict)