def collect_all_fn_performance(before_or_after):
# - fn performance on pretrain data
since = time()
dataset_name = "Pretrain"
random_subset_indices = np.random.choice(
len(pretrain_training_dataset), size=(6000), replace=False)
fn_pretrain_loader = get_loader(
dataset=Subset(pretrain_training_dataset,
random_subset_indices),
batch_size=batch_size,
drop_last=False,
)
_, _, _, fn_pretrain_preds, _, fn_pretrain_losses = \
get_fn_performance(fn_pretrain_loader, dataset_name, before_or_after)
fn_pretrain_preds = np.hstack(fn_pretrain_preds)
fn_pretrain_losses = np.vstack(fn_pretrain_losses)
print(
f"Time taken for collecting fn performance on {dataset_name} {before_or_after} updating fn weights: {time() - since} s"
)
# - fn performance on cotrain data
since = time()
dataset_name = "Cotrain_full"
_, _, _, fn_cotrain_full_preds, _, fn_cotrain_full_losses = get_fn_performance(
get_loader(cotrain_grasp_dataset,
batch_size,
shuffle=False,
drop_last=False), dataset_name, before_or_after)
fn_cotrain_full_preds = np.hstack(fn_cotrain_full_preds)
fn_cotrain_full_losses = np.vstack(fn_cotrain_full_losses)
print(
f"Time taken for collecting fitness loss on cotrain_dataset {before_or_after} updating fn weights: {time() - since} s"
)
return {
f"Pretrain_{before_or_after}":
fn_pretrain_losses.mean()
if len(fn_pretrain_losses) != 0 else 0,
f"cotrain_full_{before_or_after}":
fn_cotrain_full_losses.mean()
if len(fn_cotrain_full_losses) != 0 else 0,
f"cotrain_old_{before_or_after}":
fn_cotrain_full_losses[cotrain_old_indices].mean() if
len(fn_cotrain_full_losses[cotrain_old_indices]) != 0 else 0,
f"cotrain_new_{before_or_after}":
fn_cotrain_full_losses[cotrain_new_indices].mean() if
len(fn_cotrain_full_losses[cotrain_new_indices]) != 0 else 0,
}, {
f"FN_Preds/Pretrain_{before_or_after}":
fn_pretrain_preds if len(fn_pretrain_preds) != 0 else 0,
f"FN_Preds/cotrain_full_{before_or_after}":
fn_cotrain_full_preds
if len(fn_cotrain_full_preds) != 0 else 0,
f"FN_Preds/cotrain_old_{before_or_after}":
fn_cotrain_full_preds[cotrain_old_indices]
if len(fn_cotrain_full_preds[cotrain_old_indices]) != 0 else 0,
f"FN_Preds/cotrain_new_{before_or_after}":
fn_cotrain_full_preds[cotrain_new_indices]
if len(fn_cotrain_full_preds[cotrain_new_indices]) != 0 else 0,
}
def get_mixed_cotrain_pretrain_loader():
if hyperparams['balanced_sampling']:
return get_balanced_cotrain_loader(
cotrain_dataset=cotrain_grasp_dataset,
pretrain_dataset=pretrain_training_dataset,
batch_size=batch_size)
else:
return get_combined_loader(
get_loader(dataset=cotrain_grasp_dataset,
batch_size=int(batch_size / 2)),
get_loader(dataset=pretrain_training_dataset,
batch_size=int(batch_size / 2)))
def get_val_loader(batch_size=hyperparameters['pretrain_batch_size']):
if hyperparameters['balanced_sampling']:
batch_size = int(batch_size / 2)
return get_loader(dataset=val_dataset,
batch_size=batch_size,
collate_fn=collate_fn,
distributed=False)
def get_gn_loader():
return get_loader(object_dataset,
batch_size=batch_size,
collate_fn=lambda batch: stack(
[item['tsdf'] for item in batch]))