num_traj = config.num_traj training_reward_save = 0 for epoch in range(config.epochs): train_loss = 0 total_reward = 0 imitate_net.epsilon = 1 imitate_net.train() for batch_idx in range(config.train_size // (config.batch_size)): optimizer.zero_grad() loss_sum = Variable(torch.zeros(1)).cuda().data Rs = np.zeros((config.batch_size, 1)) for _ in range(num_traj): labels = np.zeros((config.batch_size, max_len), dtype=np.int32) data_ = next(train_gen) one_hot_labels = prepare_input_op(labels, len(unique_draw)) one_hot_labels = Variable(torch.from_numpy(one_hot_labels)).cuda() data = Variable(torch.from_numpy(data_), volatile=False).cuda() outputs, samples = imitate_net([data, one_hot_labels, max_len]) R = reinforce.generate_rewards(samples, data_, time_steps=max_len, stack_size=max_len // 2 + 1, reward=reward, power=power) R = R[0] loss = reinforce.pg_loss_var(R, samples, outputs) / num_traj loss.backward() if reward == "chamfer": Rs = Rs + R
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], jitter_program=jit) for k in dataset_sizes.keys(): test_batch_size = config.batch_size for _ in range(dataset_sizes[k][1] // test_batch_size): with torch.no_grad(): data_, labels = next(test_gen_objs[k]) labels_cont = torch.from_numpy( labels_to_cont(labels, generator.unique_draw)).cuda().float() one_hot_labels = prepare_input_op(labels, len(generator.unique_draw)) one_hot_labels = Variable( torch.from_numpy(one_hot_labels)).cuda() data = Variable(torch.from_numpy(data_)).cuda() labels = Variable(torch.from_numpy(labels)).cuda() test_output = imitate_net.test([data, one_hot_labels, max_len]) # test_output = torch.stack(test_output).permute(1, 0, 2) # acc += float((torch.argmax(test_output[:, :, :8], dim=2)[:, :(k)] == labels_cont[:, 1:-1, 0]).float().sum()) \ # / (len(labels_cont) * (k)) / len(dataset_sizes) / (dataset_sizes[k][1] // test_batch_size) acc += float((torch.argmax(torch.stack(test_output), dim=2)[:k].permute(1, 0) == labels[:, :-1]).float().sum()) \ / (len(labels) * (k)) / len(dataset_sizes) / (dataset_sizes[k][1] // test_batch_size) pred_images, correct_prog, pred_prog = parser.get_final_canvas( test_output, if_just_expressions=False, if_pred_images=True) target_images = data_[-1, :, 0, :, :].astype(dtype=bool)
def train_inference(imitate_net, path, max_epochs=None, self_training=False, ab=None): if max_epochs is None: epochs = 1000 else: epochs = max_epochs config = read_config.Config("config_synthetic.yml") if ab is not None: train_size = inference_train_size * ab else: train_size = inference_train_size generator = WakeSleepGen(f"{path}/", batch_size=config.batch_size, train_size=train_size, canvas_shape=config.canvas_shape, max_len=max_len, self_training=True) train_gen = generator.get_train_data() cad_generator = Generator() val_gen = cad_generator.val_gen(batch_size=config.batch_size, path="data/cad/cad.h5", if_augment=False) for parameter in imitate_net.encoder.parameters(): parameter.requires_grad = False optimizer = optim.Adam( [para for para in imitate_net.parameters() if para.requires_grad], weight_decay=config.weight_decay, lr=config.lr) reduce_plat = LearningRate(optimizer, init_lr=config.lr, lr_dacay_fact=0.2, patience=config.patience) best_test_loss = 1e20 torch.save(imitate_net.state_dict(), f"{path}/best_dict.pth") best_test_cd = 1e20 patience = 20 num_worse = 0 for epoch in range(epochs): start = time.time() train_loss = 0 imitate_net.train() for batch_idx in range(train_size // (config.batch_size * config.num_traj)): optimizer.zero_grad() loss = 0 # acc = 0 for _ in range(config.num_traj): data, labels = next(train_gen) # data = data[:, :, 0:1, :, :] one_hot_labels = prepare_input_op(labels, len(generator.unique_draw)) one_hot_labels = torch.from_numpy(one_hot_labels).to(device) data = data.to(device) labels = labels.to(device) outputs = imitate_net([data, one_hot_labels, max_len]) # acc += float((torch.argmax(outputs, dim=2).permute(1, 0) == labels).float().sum()) \ # / (labels.shape[0] * labels.shape[1]) / config.num_traj loss_k = ( (losses_joint(outputs, labels, time_steps=max_len + 1) / (max_len + 1)) / config.num_traj) loss_k.backward() loss += float(loss_k) del loss_k optimizer.step() train_loss += loss print(f"batch {batch_idx} train loss: {loss}") # print(f"acc: {acc}") mean_train_loss = train_loss / (train_size // (config.batch_size)) print(f"epoch {epoch} mean train loss: {mean_train_loss}") imitate_net.eval() loss = 0 # acc = 0 metrics = {"cos": 0, "iou": 0, "cd": 0} # IOU = 0 # COS = 0 CD = 0 # correct_programs = 0 # pred_programs = 0 for batch_idx in range(inference_test_size // config.batch_size): parser = ParseModelOutput(generator.unique_draw, max_len // 2 + 1, max_len, config.canvas_shape) with torch.no_grad(): labels = np.zeros((config.batch_size, max_len), dtype=np.int32) data_ = next(val_gen) one_hot_labels = prepare_input_op(labels, len(generator.unique_draw)) one_hot_labels = torch.from_numpy(one_hot_labels).cuda() data = torch.from_numpy(data_).cuda() # outputs = imitate_net([data, one_hot_labels, max_len]) # loss_k = (losses_joint(outputs, labels, time_steps=max_len + 1) / # (max_len + 1)) # loss += float(loss_k) test_outputs = imitate_net.test( [data[-1, :, 0, :, :], one_hot_labels, max_len]) # acc += float((torch.argmax(torch.stack(test_outputs), dim=2).permute(1, 0) == labels[:, :-1]).float().sum()) \ # / (len(labels) * (max_len+1)) / (inference_test_size // config.batch_size) pred_images, correct_prog, pred_prog = parser.get_final_canvas( test_outputs, if_just_expressions=False, if_pred_images=True) # correct_programs += len(correct_prog) # pred_programs += len(pred_prog) target_images = data_[-1, :, 0, :, :].astype(dtype=bool) # iou = np.sum(np.logical_and(target_images, pred_images), # (1, 2)) / \ # np.sum(np.logical_or(target_images, pred_images), # (1, 2)) # cos = cosine_similarity(target_images, pred_images) CD += np.sum(chamfer(target_images, pred_images)) # IOU += np.sum(iou) # COS += np.sum(cos) # metrics["iou"] = IOU / inference_test_size # metrics["cos"] = COS / inference_test_size metrics["cd"] = CD / inference_test_size test_losses = loss test_loss = test_losses / (inference_test_size // (config.batch_size)) if metrics["cd"] >= best_test_cd: num_worse += 1 else: num_worse = 0 best_test_cd = metrics["cd"] torch.save(imitate_net.state_dict(), f"{path}/best_dict.pth") if num_worse >= patience: # load the best model and stop training imitate_net.load_state_dict(torch.load(f"{path}/best_dict.pth")) return epoch + 1 # reduce_plat.reduce_on_plateu(metrics["cd"]) print( f"Epoch {epoch}/100 => train_loss: {mean_train_loss}, iou: {0}, cd: {metrics['cd']}, test_mse: {test_loss}, test_acc: {0}" ) # print(f"CORRECT PROGRAMS: {correct_programs}") # print(f"PREDICTED PROGRAMS: {pred_programs}") # print(f"RATIO: {correct_programs/pred_programs}") end = time.time() print(f"Inference train time {end-start}") del test_losses, outputs, test_outputs return epochs
def infer_programs(imitate_net, self_training=False, ab=None): config = read_config.Config("config_cad.yml") # 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] config.train_size = 10000 config.test_size = 3000 imitate_net.eval() imitate_net.epsilon = 0 parser = ParseModelOutput(unique_draw, max_len // 2 + 1, max_len, config.canvas_shape) generator = Generator() test_gen = generator.test_gen(batch_size=config.batch_size, path="data/cad/cad.h5", if_augment=False) pred_expressions = [] Rs = 0 CDs = 0 Target_images = [] for batch_idx in range(config.test_size // config.batch_size): with torch.no_grad(): print(f"Inferring test cad batch: {batch_idx}") data_ = next(test_gen) labels = np.zeros((config.batch_size, max_len), dtype=np.int32) one_hot_labels = prepare_input_op(labels, len(unique_draw)) one_hot_labels = torch.from_numpy(one_hot_labels).to(device) data = torch.from_numpy(data_).to(device) all_beams, next_beams_prob, all_inputs = imitate_net.beam_search( [data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len) beam_labels = beams_parser(all_beams, data_.shape[1], beam_width=beam_width) beam_labels_numpy = np.zeros( (config.batch_size * beam_width, max_len), dtype=np.int32) Target_images.append(data_[-1, :, 0, :, :]) for i in range(data_.shape[1]): beam_labels_numpy[i * beam_width:(i + 1) * beam_width, :] = beam_labels[i] # find expression from these predicted beam labels expressions = [""] * config.batch_size * beam_width for i in range(config.batch_size * beam_width): for j in range(max_len): expressions[i] += unique_draw[beam_labels_numpy[i, j]] for index, prog in enumerate(expressions): expressions[index] = prog.split("$")[0] pred_expressions += expressions predicted_images = image_from_expressions(parser, expressions) target_images = data_[-1, :, 0, :, :].astype(dtype=bool) target_images_new = np.repeat(target_images, axis=0, repeats=beam_width) beam_CD = chamfer(target_images_new, predicted_images) CD = np.zeros((config.batch_size, 1)) for r in range(config.batch_size): CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width]) CDs += np.mean(CD) for j in range(0, config.batch_size): f, a = plt.subplots(1, beam_width + 1, figsize=(30, 3)) a[0].imshow(data_[-1, j, 0, :, :], cmap="Greys_r") a[0].axis("off") a[0].set_title("target") for i in range(1, beam_width + 1): a[i].imshow(predicted_images[j * beam_width + i - 1], cmap="Greys_r") a[i].set_title("{}".format(i)) a[i].axis("off") plt.savefig("best_lest/" + "{}.png".format(batch_idx * config.batch_size + j), transparent=0) plt.close("all") # with open("best_st_expressions.txt", "w") as file: # for e in pred_expressions: # file.write(f"{e}\n") # break return CDs / (config.test_size // config.batch_size)
def train_inference(inference_net, iter): config = read_config.Config("config_synthetic.yml") generator = WakeSleepGen( f"wake_sleep_data/inference/{iter}/labels/labels.pt", f"wake_sleep_data/inference/{iter}/labels/val/labels.pt", batch_size=config.batch_size, train_size=inference_train_size, test_size=inference_test_size, canvas_shape=config.canvas_shape, max_len=max_len) train_gen = generator.get_train_data() test_gen = generator.get_test_data() encoder_net, imitate_net = inference_net optimizer = optim.Adam( [para for para in imitate_net.parameters() if para.requires_grad], weight_decay=config.weight_decay, lr=config.lr) reduce_plat = LearningRate(optimizer, init_lr=config.lr, lr_dacay_fact=0.2, patience=config.patience) best_test_loss = 1e20 best_imitate_dict = imitate_net.state_dict() prev_test_cd = 1e20 prev_test_iou = 0 patience = 5 num_worse = 0 for epoch in range(50): train_loss = 0 Accuracies = [] imitate_net.train() for batch_idx in range(inference_train_size // (config.batch_size * config.num_traj)): optimizer.zero_grad() loss = Variable(torch.zeros(1)).to(device).data for _ in range(config.num_traj): batch_data, batch_labels = next(train_gen) batch_data = batch_data.to(device) batch_labels = batch_labels.to(device) batch_data = batch_data[:, :, 0:1, :, :] one_hot_labels = prepare_input_op(batch_labels, vocab_size) one_hot_labels = Variable( torch.from_numpy(one_hot_labels)).to(device) outputs = imitate_net([batch_data, one_hot_labels, max_len]) loss_k = (losses_joint( outputs, batch_labels, time_steps=max_len + 1) / (max_len + 1)) / config.num_traj loss_k.backward() loss += loss_k.data del loss_k optimizer.step() train_loss += loss print(f"batch {batch_idx} train loss: {loss.cpu().numpy()}") mean_train_loss = train_loss / (inference_train_size // (config.batch_size)) print( f"epoch {epoch} mean train loss: {mean_train_loss.cpu().numpy()}") imitate_net.eval() loss = Variable(torch.zeros(1)).to(device) metrics = {"cos": 0, "iou": 0, "cd": 0} IOU = 0 COS = 0 CD = 0 for batch_idx in range(inference_test_size // config.batch_size): with torch.no_grad(): batch_data, batch_labels = next(test_gen) batch_data = batch_data.to(device) batch_labels = batch_labels.to(device) one_hot_labels = prepare_input_op(batch_labels, vocab_size) one_hot_labels = Variable( torch.from_numpy(one_hot_labels)).to(device) test_outputs = imitate_net( [batch_data, one_hot_labels, max_len]) loss += (losses_joint( test_outputs, batch_labels, time_steps=max_len + 1) / (max_len + 1)) test_output = imitate_net.test( [batch_data, one_hot_labels, max_len]) pred_images, correct_prog, pred_prog = generator.parser.get_final_canvas( test_output, if_just_expressions=False, if_pred_images=True) target_images = batch_data.cpu().numpy()[-1, :, 0, :, :].astype( dtype=bool) iou = np.sum(np.logical_and(target_images, pred_images), (1, 2)) / \ np.sum(np.logical_or(target_images, pred_images), (1, 2)) cos = cosine_similarity(target_images, pred_images) CD += np.sum(chamfer(target_images, pred_images)) IOU += np.sum(iou) COS += np.sum(cos) metrics["iou"] = IOU / inference_test_size metrics["cos"] = COS / inference_test_size metrics["cd"] = CD / inference_test_size test_losses = loss.data test_loss = test_losses.cpu().numpy() / (inference_test_size // (config.batch_size)) if test_loss >= best_test_loss: num_worse += 1 else: num_worse = 0 best_test_loss = test_loss best_imitate_dict = imitate_net.state_dict() if num_worse >= patience: # load the best model and stop training imitate_net.load_state_dict(best_imitate_dict) break reduce_plat.reduce_on_plateu(metrics["cd"]) print("Epoch {}/{}=> train_loss: {}, iou: {}, cd: {}, test_mse: {}". format( epoch, config.epochs, mean_train_loss.cpu().numpy(), metrics["iou"], metrics["cd"], test_loss, )) print(f"CORRECT PROGRAMS: {len(generator.correct_programs)}") del test_losses, test_outputs
def infer_programs(imitate_net, path, self_training=False, ab=None): save_viz = False config = read_config.Config("config_cad.yml") # 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] config.train_size = 10000 config.test_size = 3000 imitate_net.eval() imitate_net.epsilon = 0 parser = ParseModelOutput(unique_draw, max_len // 2 + 1, max_len, config.canvas_shape) pred_expressions = [] if ab is not None: pred_labels = np.zeros((config.train_size * ab, max_len)) else: pred_labels = np.zeros((config.train_size, max_len)) image_path = f"{path}/images/" results_path = f"{path}/results/" labels_path = f"{path}/labels/" os.makedirs(os.path.dirname(image_path), exist_ok=True) os.makedirs(os.path.dirname(results_path), exist_ok=True) os.makedirs(os.path.dirname(labels_path), exist_ok=True) os.makedirs(os.path.dirname(labels_path + "val/"), exist_ok=True) generator = Generator() train_gen = generator.train_gen(batch_size=config.batch_size, path="data/cad/cad.h5", if_augment=False) val_gen = generator.val_gen(batch_size=config.batch_size, path="data/cad/cad.h5", if_augment=False) Rs = 0 CDs = 0 Target_images = [] start = time.time() pred_images = np.zeros((config.train_size, 64, 64)) for batch_idx in range(config.train_size // config.batch_size): with torch.no_grad(): print(f"Inferring cad batch: {batch_idx}") data_ = next(train_gen) labels = np.zeros((config.batch_size, max_len), dtype=np.int32) one_hot_labels = prepare_input_op(labels, len(unique_draw)) one_hot_labels = torch.from_numpy(one_hot_labels).to(device) data = torch.from_numpy(data_).to(device) all_beams, next_beams_prob, all_inputs = imitate_net.beam_search( [data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len) beam_labels = beams_parser(all_beams, data_.shape[1], beam_width=beam_width) beam_labels_numpy = np.zeros( (config.batch_size * beam_width, max_len), dtype=np.int32) Target_images.append(data_[-1, :, 0, :, :]) for i in range(data_.shape[1]): beam_labels_numpy[i * beam_width:(i + 1) * beam_width, :] = beam_labels[i] # find expression from these predicted beam labels expressions = [""] * config.batch_size * beam_width for i in range(config.batch_size * beam_width): for j in range(max_len): expressions[i] += unique_draw[beam_labels_numpy[i, j]] for index, prog in enumerate(expressions): expressions[index] = prog.split("$")[0] pred_expressions += expressions predicted_images = image_from_expressions(parser, expressions) target_images = data_[-1, :, 0, :, :].astype(dtype=bool) target_images_new = np.repeat(target_images, axis=0, repeats=beam_width) # beam_R = np.sum(np.logical_and(target_images_new, predicted_images), # (1, 2)) / np.sum(np.logical_or(target_images_new, predicted_images), (1, 2)) # # R = np.zeros((config.batch_size, 1)) # for r in range(config.batch_size): # R[r, 0] = max(beam_R[r * beam_width:(r + 1) * beam_width]) # # Rs += np.mean(R) beam_CD = chamfer(target_images_new, predicted_images) # select best expression by chamfer distance if ab is None: best_labels = np.zeros((config.batch_size, max_len)) for r in range(config.batch_size): idx = np.argmin(beam_CD[r * beam_width:(r + 1) * beam_width]) best_labels[r] = beam_labels[r][idx] pred_labels[batch_idx * config.batch_size:batch_idx * config.batch_size + config.batch_size] = best_labels else: best_labels = np.zeros((config.batch_size * ab, max_len)) for r in range(config.batch_size): sorted_idx = np.argsort(beam_CD[r * beam_width:(r + 1) * beam_width])[:ab] best_labels[r * ab:r * ab + ab] = beam_labels[r][sorted_idx] pred_labels[batch_idx * config.batch_size * ab:batch_idx * config.batch_size * ab + config.batch_size * ab] = best_labels CD = np.zeros((config.batch_size, 1)) for r in range(config.batch_size): CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width]) pred_images[batch_idx * config.batch_size + r] = predicted_images[r * beam_width + np.argmin( beam_CD[r * beam_width:(r + 1) * beam_width])] CDs += np.mean(CD) if save_viz: for j in range(0, config.batch_size): f, a = plt.subplots(1, beam_width + 1, figsize=(30, 3)) a[0].imshow(data_[-1, j, 0, :, :], cmap="Greys_r") a[0].axis("off") a[0].set_title("target") for i in range(1, beam_width + 1): a[i].imshow(predicted_images[j * beam_width + i - 1], cmap="Greys_r") a[i].set_title("{}".format(i)) a[i].axis("off") plt.savefig( image_path + "{}.png".format(batch_idx * config.batch_size + j), transparent=0) plt.close("all") save_viz = False print("Inferring cad average chamfer distance: {}".format( CDs / (config.train_size // config.batch_size)), flush=True) Rs = Rs / (config.train_size // config.batch_size) CDs = CDs / (config.train_size // config.batch_size) print(Rs, CDs) results = {"iou": Rs, "chamferdistance": CDs} with open(results_path + "results_beam_width_{}.org".format(beam_width), 'w') as outfile: json.dump(results, outfile) torch.save(pred_labels, labels_path + "labels.pt") # torch.save(pred_images, labels_path + "images.pt") if self_training: if ab is None: torch.save(np.concatenate(Target_images, axis=0), labels_path + "images.pt") else: torch.save( np.repeat(np.concatenate(Target_images, axis=0), ab, axis=0), labels_path + "images.pt") test_gen = generator.test_gen(batch_size=config.batch_size, path="data/cad/cad.h5", if_augment=False) pred_expressions = [] Rs = 0 CDs = 0 Target_images = [] for batch_idx in range(config.test_size // config.batch_size): with torch.no_grad(): print(f"Inferring test cad batch: {batch_idx}") data_ = next(test_gen) labels = np.zeros((config.batch_size, max_len), dtype=np.int32) one_hot_labels = prepare_input_op(labels, len(unique_draw)) one_hot_labels = torch.from_numpy(one_hot_labels).to(device) data = torch.from_numpy(data_).to(device) all_beams, next_beams_prob, all_inputs = imitate_net.beam_search( [data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len) beam_labels = beams_parser(all_beams, data_.shape[1], beam_width=beam_width) beam_labels_numpy = np.zeros( (config.batch_size * beam_width, max_len), dtype=np.int32) Target_images.append(data_[-1, :, 0, :, :]) for i in range(data_.shape[1]): beam_labels_numpy[i * beam_width:(i + 1) * beam_width, :] = beam_labels[i] # find expression from these predicted beam labels expressions = [""] * config.batch_size * beam_width for i in range(config.batch_size * beam_width): for j in range(max_len): expressions[i] += unique_draw[beam_labels_numpy[i, j]] for index, prog in enumerate(expressions): expressions[index] = prog.split("$")[0] pred_expressions += expressions predicted_images = image_from_expressions(parser, expressions) target_images = data_[-1, :, 0, :, :].astype(dtype=bool) target_images_new = np.repeat(target_images, axis=0, repeats=beam_width) beam_CD = chamfer(target_images_new, predicted_images) CD = np.zeros((config.batch_size, 1)) for r in range(config.batch_size): CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width]) CDs += np.mean(CD) print(f"TEST CD: {CDs / (config.test_size // config.batch_size)}") end = time.time() print(f"Inference time: {end-start}")