data = data.permute(1, 0, 2, 3, 4, 5)
labels = Variable(torch.from_numpy(labels)).cuda()
test_output = imitate_net([data, one_hot_labels, k])
l = losses_joint(test_output, labels, time_steps=k + 1).data
test_losses += l
if cuda_devices > 1:
test_output = imitate_net.module.test(
[data, one_hot_labels, k])
else:
test_output = imitate_net.test([data, one_hot_labels, k])
stack, _, _ = parser.get_final_canvas(test_output,
if_pred_images=True,
if_just_expressions=False)
data_ = data_[-1, :, 0, :, :, :]
R = np.sum(np.logical_and(stack, data_),
(1, 2, 3)) / (np.sum(np.logical_or(stack, data_),
(1, 2, 3)) + 1)
test_reward += np.sum(R)
test_reward = test_reward / (test_size // batch_size) / (
(batch_size // types_prog) * types_prog)
test_loss = test_losses.cpu().numpy() / (config.test_size //
config.batch_size) / types_prog
log_value('test_loss', test_loss, epoch)
log_value('test_IOU',
test_reward / (config.test_size // config.batch_size), epoch)
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
with torch.no_grad():
data_, labels = next(test_gen_objs[k])
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_outputs = imitate_net([data, one_hot_labels, k])
#loss += (losses_joint(test_outputs, labels, time_steps=k + 1) /
# (k + 1)) / types_prog
test_output = imitate_net.test([data, one_hot_labels, max_len])
#acc += float((torch.argmax(torch.stack(test_output), dim=2)[:k].permute(1, 0) == labels[:, :-1]).float().sum()) \
# / (len(labels) * (k+1)) / types_prog / (config.test_size // config.batch_size)
pred_images, correct_prog, pred_prog = parser.get_final_canvas(
test_output,
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))
#beam_CD += get_cd(imitate_net, data, one_hot_labels, k)
#IOU += np.sum(iou)
#COS += np.sum(cos)
#metrics["iou"] = IOU / config.test_size
data_ = data_[:, :, 0:config.top_k + 1, :, :, :]
one_hot_labels = prepare_input_op(labels, len(generator.unique_draw))
one_hot_labels = Variable(torch.from_numpy(one_hot_labels),
volatile=True).cuda()
data = Variable(torch.from_numpy(data_)).cuda()
labels = Variable(torch.from_numpy(labels)).cuda()
# This is for data parallelism purpose
data = data.permute(1, 0, 2, 3, 4, 5)
if cuda_devices > 1:
outputs = imitate_net.module.test([data, one_hot_labels, max_len])
else:
outputs = imitate_net.test([data, one_hot_labels, max_len])
stack, _, expressions = parser.get_final_canvas(
outputs, if_pred_images=True, if_just_expressions=False)
Predicted_expressions += expressions
target_expressions = parser.labels2exps(labels, k)
Target_expressions += target_expressions
# stacks = parser.expression2stack(expressions)
data_ = data_[-1, :, 0, :, :, :]
R = np.sum(np.logical_and(stack, data_),
(1, 2, 3)) / (np.sum(np.logical_or(stack, data_),
(1, 2, 3)) + 1)
Rs += np.sum(R)
total_iou += Rs
IOU[k] = Rs / (
(dataset_sizes[k][1] // config.batch_size) * config.batch_size)
print("IOU for {} len program: ".format(k), IOU[k])
total_iou = total_iou / config.test_size
def infer_progams(csgnet, train_dataset, val_dataset):
parser = ParseModelOutput(unique_draws,
max_len // 2 + 1,
max_len, [64, 64, 64],
primitives=primitives)
csgnet.eval()
datasets = [("train", train_dataset), ("val", val_dataset)]
dataset_expressions = []
dataset_stacks = []
dataset_labels = []
for name, dataset in datasets:
predicted_expressions = []
predicted_stacks = []
predicted_labels = []
IOU = {}
total_iou = 0
Rs = 0.0
batch_idx = 0
count = 0
for batch in dataset:
with torch.no_grad():
print(f"batch {batch_idx}/{len(train_dataset)}")
batch_idx += 1
count += len(batch)
vis_voxels(batch.squeeze().numpy()[:5], "gt")
batch = batch.to(device)
outputs = csgnet.test2(batch, max_len)
labels = [
torch.max(o, 1)[1].data.cpu().numpy() for o in outputs
]
labels += [np.full((len(batch), ), len(unique_draws) - 1)]
stack, _, expressions = parser.get_final_canvas(
outputs, if_pred_images=True, if_just_expressions=False)
vis_voxels(stack[:5], "gen")
break
predicted_expressions += expressions
predicted_stacks.append(stack)
predicted_labels.append(np.stack(labels).transpose())
# stacks = parser.expression2stack(expressions)
data_ = batch.squeeze().cpu().numpy()
R = np.sum(np.logical_and(stack, data_),
(1, 2, 3)) / (np.sum(np.logical_or(stack, data_),
(1, 2, 3)) + 1)
Rs += np.sum(R)
IOU = Rs / count
print(f"IOU on ShapeNet {name}: {IOU}")
dataset_expressions.append(predicted_expressions)
dataset_stacks.append(np.concatenate(predicted_stacks, axis=0))
dataset_labels.append(np.concatenate(predicted_labels, axis=0))
train_samples = list(zip(dataset_labels[0], list(dataset_stacks[0])))
val_samples = list(zip(dataset_labels[1], list(dataset_stacks[1])))
train_dataset = DataLoader(train_samples,
batch_size=config.batch_size,
shuffle=True,
collate_fn=_col)
val_dataset = DataLoader(val_samples,
batch_size=config.batch_size,
shuffle=False,
collate_fn=_col)
return train_dataset, val_dataset
def train_model(csgnet, train_dataset, val_dataset, max_epochs=None):
if max_epochs is None:
epochs = 100
else:
epochs = max_epochs
optimizer = optim.Adam(
[para for para in csgnet.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,
lr_decay_epoch=3,
patience=config.patience)
best_state_dict = None
patience = 3
prev_test_loss = 1e20
prev_test_reward = 0
num_worse = 0
for epoch in range(100):
train_loss = 0
Accuracies = []
csgnet.train()
# Number of times to accumulate gradients
num_accums = config.num_traj
batch_idx = 0
count = 0
for batch in train_dataset:
labels = np.stack([x[0] for x in batch])
data = np.stack([x[1] for x in batch])
if not len(labels) == config.batch_size:
continue
optimizer.zero_grad()
loss_sum = Variable(torch.zeros(1)).cuda().data
one_hot_labels = prepare_input_op(labels, len(unique_draws))
one_hot_labels = Variable(torch.from_numpy(one_hot_labels)).cuda()
data = Variable(
torch.from_numpy(data)).cuda().unsqueeze(-1).float()
labels = Variable(torch.from_numpy(labels)).cuda()
# forward pass
outputs = csgnet.forward2([data, one_hot_labels, max_len])
loss = losses_joint(outputs, labels,
time_steps=max_len + 1) / num_accums
loss.backward()
loss_sum += loss.data
batch_idx += 1
count += len(data)
if batch_idx % num_accums == 0:
# Clip the gradient to fixed value to stabilize training.
torch.nn.utils.clip_grad_norm_(csgnet.parameters(), 20)
optimizer.step()
l = loss_sum
train_loss += l
# print(f'train loss batch {batch_idx}: {l}')
mean_train_loss = (train_loss * num_accums) / (count //
config.batch_size)
print(f'train loss epoch {epoch}: {float(mean_train_loss)}')
del data, loss, loss_sum, train_loss, outputs
test_losses = 0
acc = 0
csgnet.eval()
test_reward = 0
batch_idx = 0
count = 0
for batch in val_dataset:
labels = np.stack([x[0] for x in batch])
data = np.stack([x[1] for x in batch])
if not len(labels) == config.batch_size:
continue
parser = ParseModelOutput(unique_draws,
stack_size=(max_len + 1) // 2 + 1,
steps=max_len,
canvas_shape=[64, 64, 64],
primitives=primitives)
with torch.no_grad():
one_hot_labels = prepare_input_op(labels, len(unique_draws))
one_hot_labels = Variable(
torch.from_numpy(one_hot_labels)).cuda()
data = Variable(
torch.from_numpy(data)).cuda().unsqueeze(-1).float()
labels = Variable(torch.from_numpy(labels)).cuda()
test_output = csgnet.forward2([data, one_hot_labels, max_len])
l = losses_joint(test_output, labels,
time_steps=max_len + 1).data
test_losses += l
acc += float((torch.argmax(torch.stack(test_output), dim=2).permute(1, 0) == labels).float().sum()) \
/ (labels.shape[0] * labels.shape[1])
test_output = csgnet.test2(data, max_len)
stack, _, _ = parser.get_final_canvas(
test_output,
if_pred_images=True,
if_just_expressions=False)
data_ = data.squeeze().cpu().numpy()
R = np.sum(np.logical_and(stack, data_),
(1, 2, 3)) / (np.sum(np.logical_or(stack, data_),
(1, 2, 3)) + 1)
test_reward += np.sum(R)
batch_idx += 1
count += len(data)
test_reward = test_reward / count
test_loss = test_losses / (count // config.batch_size)
acc = acc / (count // config.batch_size)
if test_loss < prev_test_loss:
prev_test_loss = test_loss
best_state_dict = csgnet.state_dict()
num_worse = 0
else:
num_worse += 1
if num_worse >= patience:
csgnet.load_state_dict(best_state_dict)
break
print(f'test loss epoch {epoch}: {float(test_loss)}')
print(f'test IOU epoch {epoch}: {test_reward}')
print(f'test acc epoch {epoch}: {acc}')
if config.if_schedule:
reduce_plat.reduce_on_plateu(-test_reward)
del test_losses, test_output
if test_reward > prev_test_reward:
prev_test_reward = test_reward