def per_condition_loss(d, c, conds, model, args, idx=1):
tmp1, tmp2 = torch.split(d, int(d.size()[-1] / 2), dim=1)
condition_tensor = d.clone()
tmp1, tmp2 = torch.split(condition_tensor,
int(condition_tensor.size()[-1] / 2),
dim=1)
for kk in conds:
tmp1[:, kk], tmp2[:, kk] = 0, 0
cond_d = torch.cat((tmp1, tmp2), 1)
# Run 10 times for resampling
my_recon_list, my_z_means_list, my_log_var_list = [], [], []
for resample in range(10):
recon_batch, z_means, log_var = model(c.clone().cuda(args.gpu_id),
cond_d.clone().cuda(args.gpu_id))
my_recon_list.append(recon_batch)
my_z_means_list.append(z_means)
my_log_var_list.append(log_var)
recon_batch = torch.mean(torch.stack(my_recon_list), dim=0)
z_means = torch.mean(torch.stack(my_z_means_list), dim=0)
log_var = torch.mean(torch.stack(my_log_var_list), dim=0)
loss_fn = str_to_object(args.loss_fn)
ELBO_loss, RCL_loss, KLD_loss, _, _ = loss_fn(
c.cuda(args.gpu_id), recon_batch.cuda(args.gpu_id), z_means, log_var,
args)
if idx == 1:
return ELBO_loss.item()
elif idx == 5:
return ELBO_loss.item(), RCL_loss.item(), KLD_loss.item()
else:
return recon_batch, z_means, log_var
def make_dataframe(kl_per_lt, kl_all_lt, selected_features, feature_names,
z_means, c, recon_batch, log_var, conds, popped_cond, args):
z_means_x, z_var_x, z_means_y, z_var_y = [], [], [], []
all_kl, all_lt = [], []
loss_fn = str_to_object(args.loss_fn)
total_ELBO, total_RCL, total_KLD, _, _ = loss_fn(
c.cuda(args.gpu_id), recon_batch.cuda(args.gpu_id), z_means, log_var,
args)
selected_features["ELBO"].append(total_ELBO.item())
selected_features["RCL"].append(total_RCL.item())
selected_features["KLD"].append(total_KLD.item())
if popped_cond != []:
selected_features["selected_feature_number"].append(str(popped_cond))
if args.data_type == 'aics_features':
name = feature_names[popped_cond]
selected_features["selected_feature_name"].append(name)
else:
selected_features["selected_feature_name"].append(None)
else:
selected_features["selected_feature_number"].append(None)
selected_features["selected_feature_name"].append(None)
for ii in range(z_means.size()[-1]):
_, rcl_per_lt_temp, kl_per_lt_temp, _, _ = loss_fn(
c.cuda(args.gpu_id), recon_batch.cuda(args.gpu_id), z_means[:, ii],
log_var[:, ii], args)
all_kl = np.append(all_kl, kl_per_lt_temp.item())
all_lt.append(ii)
# print('greedy encoding plots', c.size()[-1] - len(conds))
kl_per_lt["num_conds"].append(c.size()[-1] - len(conds))
if popped_cond != []:
kl_per_lt["popped_cond"].append(popped_cond)
else:
kl_per_lt["popped_cond"].append(None)
kl_per_lt["latent_dim"].append(ii)
kl_per_lt["kl_divergence"].append(kl_per_lt_temp.item())
kl_per_lt['cond_comb'].append(str([i for i in conds]))
all_kl, all_lt = list(zip(*sorted(zip(all_kl, all_lt))))
all_kl = list(all_kl)
all_lt = list(all_lt)
z_means_x = np.append(z_means_x, z_means[:, all_lt[-1]].data.cpu().numpy())
z_means_y = np.append(z_means_y, z_means[:, all_lt[-2]].data.cpu().numpy())
z_var_x = np.append(z_var_x, log_var[:, all_lt[-1]].data.cpu().numpy())
z_var_y = np.append(z_var_y, log_var[:, all_lt[-2]].data.cpu().numpy())
kl_all_lt['z_means_x'].append(z_means_x)
kl_all_lt['z_means_y'].append(z_means_y)
kl_all_lt['z_var_x'].append(z_var_x)
kl_all_lt['z_var_y'].append(z_var_y)
kl_all_lt['num_conds'].append(c.size()[-1] - len(conds))
kl_all_lt['cond_comb'].append(str([i for i in conds]))
return kl_per_lt, kl_all_lt, selected_features
def get_model(model_fn, model_kwargs: Optional[Dict] = None) -> nn.Module:
model_fn = str_to_object(model_fn)
try:
return model_fn(**model_kwargs)
except:
a = dict([(key, value) for key, value in model_kwargs.items()
if key != "sklearn_data" and key != 'projection_dim'
and key != 'mask_percentage'])
return model_fn(**a)
def make_plot_encoding_greedy(args: argparse.Namespace,
model,
df: pd.DataFrame,
c,
d,
feature_names=None,
save=True,
proj_matrix=None) -> None:
"""
c and d are X_test and C_test
"""
sns.set_context("talk")
path_save_dir = Path(args.path_save_dir)
vis_enc = str_to_object(
"CVAE_testbed.metrics.greedy_visualize_encoder.GreedyVisualizeEncoder")
try:
conds = [i for i in range(args.model_kwargs["dec_layers"][-1][-1])]
except:
conds = [i for i in range(args.model_kwargs["dec_layers"][-1])]
kl_per_lt, kl_all_lt, selected_features, first_features = vis_enc(
args,
model,
conds,
c[-1, :].clone(),
d[-1, :].clone(),
kl_per_lt=None,
kl_all_lt=None,
selected_features=None,
feature_names=feature_names)
kl_per_lt, kl_all_lt, selected_features, first_features = pd.DataFrame(
kl_per_lt), pd.DataFrame(kl_all_lt), pd.DataFrame(
selected_features), pd.DataFrame(first_features)
if save is True:
path_csv = path_save_dir / Path("kl_per_lt.csv")
kl_per_lt.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_csv = path_save_dir / Path("kl_all_lt.csv")
kl_all_lt.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_csv = path_save_dir / Path("selected_features.csv")
selected_features.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_csv = path_save_dir / Path("first_features.csv")
first_features.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
fig, ax1 = plt.subplots(1, 1, figsize=(7 * 1, 4))
sns.lineplot(ax=ax1,
data=kl_per_lt,
x='latent_dim',
y='kl_divergence',
estimator='mean')
if save is True:
path_save_fig = path_save_dir / Path("greedy_elbo_kld_rcl_dims.png")
fig.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
fig2, ax = plt.subplots(1, 1, figsize=(7 * 8, 4))
first_features.sort_values(by='ELBO', ascending=False, inplace=True)
if all(pd.isna(first_features['selected_feature_name'])):
bar_fig = sns.lineplot(data=first_features,
ax=ax,
x='selected_feature_number',
y='ELBO',
label='ELBO',
sort=False)
sns.scatterplot(data=first_features,
ax=ax,
x='selected_feature_number',
y='ELBO',
s=100,
color=".2")
sns.lineplot(data=first_features,
ax=ax,
x='selected_feature_number',
y='RCL',
label='RCL',
sort=False)
sns.scatterplot(data=first_features,
ax=ax,
x='selected_feature_number',
y='RCL',
s=100,
color=".2")
else:
bar_fig = sns.lineplot(data=first_features,
ax=ax,
x='selected_feature_name',
y='ELBO',
label="ELBO",
sort=False)
sns.scatterplot(data=first_features,
ax=ax,
x='selected_feature_name',
y='ELBO',
s=100,
color=".2")
sns.lineplot(data=first_features,
ax=ax,
x='selected_feature_name',
y='RCL',
label="RCL",
sort=False)
sns.scatterplot(data=first_features,
ax=ax,
x='selected_feature_name',
y='RCL',
s=100,
color=".2")
for item in bar_fig.get_xticklabels():
item.set_rotation(45)
ax.set_title('ELBO per selected first feature')
ax.set_xlabel('Selected feature')
ax.set_ylabel('ELBO')
if save is True:
path_save_fig = path_save_dir / Path(
"greedy_barplots_first_selection.png")
fig2.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
fig2, ax = plt.subplots(1, 1, figsize=(7 * 8, 4))
print(selected_features)
selected_features.sort_values(by='ELBO', ascending=False, inplace=True)
if all(pd.isna(selected_features['selected_feature_name'])):
bar_fig = sns.lineplot(data=selected_features,
ax=ax,
x='selected_feature_number',
y='ELBO',
label='ELBO',
sort=False)
sns.scatterplot(data=selected_features,
ax=ax,
x='selected_feature_number',
y='ELBO',
s=100,
color=".2")
sns.lineplot(data=selected_features,
ax=ax,
x='selected_feature_number',
y='RCL',
label='RCL',
sort=False)
sns.scatterplot(data=selected_features,
ax=ax,
x='selected_feature_number',
y='RCL',
s=100,
color=".2")
else:
bar_fig = sns.lineplot(data=selected_features,
ax=ax,
x='selected_feature_name',
y='ELBO',
label='ELBO',
sort=False)
sns.scatterplot(data=selected_features,
ax=ax,
x='selected_feature_name',
y='ELBO',
s=100,
color=".2")
sns.lineplot(data=selected_features,
ax=ax,
x='selected_feature_name',
y='RCL',
label='RCL',
sort=False)
sns.scatterplot(data=selected_features,
ax=ax,
x='selected_feature_name',
y='RCL',
s=100,
color=".2")
for item in bar_fig.get_xticklabels():
item.set_rotation(45)
ax.set_title('ELBO per selected feature')
ax.set_xlabel('Selected feature')
ax.set_ylabel('ELBO')
if save is True:
path_save_fig = path_save_dir / Path("greedy_barplots.png")
fig2.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
def test(
args,
epoch,
loss_fn,
X_test,
C_test,
Cond_indices_test,
batch_size,
model,
optimizer,
gpu_id,
model_kwargs,
):
model.eval()
test_loss, rcl_loss, kld_loss = 0, 0, 0
rcl_per_condition_loss, kld_per_condition_loss = (
torch.zeros(X_test.size()[-1] + 1),
torch.zeros(X_test.size()[-1] + 1),
)
batch_rcl, batch_kld = (
torch.empty([0]),
torch.empty([0]),
)
batch_length = 0
with torch.no_grad():
for j, i in enumerate(X_test):
optimizer.zero_grad()
c, d, cond_labels = X_test[j], C_test[j], Cond_indices_test[j]
recon_batch, mu, log_var = model(c.cuda(gpu_id), d.cuda(gpu_id))
loss_fn = str_to_object(args.loss_fn)
loss, rcl, kld, rcl_per_element, kld_per_element = loss_fn(
c.cuda(gpu_id), recon_batch.cuda(gpu_id), mu, log_var, args)
test_loss += loss.item()
rcl_loss += rcl.item()
kld_loss += kld.item()
for jj, ii in enumerate(torch.unique(cond_labels)):
this_cond_positions = cond_labels == ii
if len(torch.unique(cond_labels)) == c.size()[-1] + 1:
batch_rcl = torch.cat(
[
batch_rcl.cuda(gpu_id),
torch.sum(rcl_per_element[this_cond_positions],
dim=0).view(1, -1),
],
0,
)
batch_kld = torch.cat(
[
batch_kld.cuda(gpu_id),
torch.sum(kld_per_element[this_cond_positions],
dim=0).view(1, -1),
],
0,
)
batch_length += 1
this_cond_rcl = torch.sum(rcl_per_element[this_cond_positions])
this_cond_kld = torch.sum(kld_per_element[this_cond_positions])
rcl_per_condition_loss[jj] += this_cond_rcl.item()
kld_per_condition_loss[jj] += this_cond_kld.item()
num_batches = len(X_test)
print("====> Epoch: {} Test losses: {:.4f}".format(epoch, test_loss /
num_batches))
print("====> RCL loss: {:.4f}".format(rcl_loss / num_batches))
print("====> KLD loss: {:.4f}".format(kld_loss / num_batches))
batch_rcl, batch_kld = batch_rcl, batch_kld
return (
test_loss / num_batches,
rcl_loss / num_batches,
kld_loss / num_batches,
rcl_per_condition_loss / num_batches,
kld_per_condition_loss / num_batches,
batch_rcl,
batch_kld,
batch_length,
)
def train_model():
"""
Trains a model
"""
tic = time.time()
args = get_args()
path_save_dir = Path(args.path_save_dir)
if path_save_dir.exists():
raise ValueError(f"Save directory already exists! ({path_save_dir})")
path_save_dir.mkdir(parents=True)
logging.getLogger().addHandler(
logging.FileHandler(path_save_dir / Path("run.log"), mode="w"))
save_args(args, path_save_dir / Path("training_options.json"))
device = (torch.device("cuda", args.gpu_id)
if torch.cuda.is_available() else torch.device("cpu"))
LOGGER.info(f"Using device: {device}")
feature_names = None
proj_matrix = None
if args.data_type == "mnist":
load_data = str_to_object(args.dataloader)
train_iterator, test_iterator = load_data(args.batch_size,
args.model_kwargs)
elif args.data_type == "aics_features":
load_data = str_to_object(args.dataloader)
test_instance = load_data(
args.num_batches,
args.batch_size,
args.model_kwargs,
corr=False,
train=True,
mask=False,
)
X_train, C_train, Cond_indices_train = test_instance.get_train_data()
X_test, C_test, Cond_indices_test = test_instance.get_test_data()
feature_names = test_instance.get_feature_names()
this_dataloader_color = test_instance.get_color()
# print('CVAE train')
# print(X_train.size(), X_test.size())
elif args.data_type == "synthetic":
if "mask_percentage" in args.model_kwargs:
mask_bool = True
else:
mask_bool = False
load_data = str_to_object(args.dataloader)
if "projection_dim" in args.model_kwargs:
X_train, C_train, Cond_indices_train, proj_matrix = load_data(
args.num_batches,
args.batch_size,
args.model_kwargs,
corr=False,
train=True,
mask=mask_bool,
).get_all_items()
path_csv = path_save_dir / Path("projection_options.pt")
print(proj_matrix)
with path_csv.open("wb") as fo:
torch.save(proj_matrix, fo)
LOGGER.info(f"Saved: {path_csv}")
test_instance = load_data(
args.num_batches,
args.batch_size,
args.model_kwargs,
corr=False,
train=False,
P=proj_matrix,
mask=mask_bool,
)
X_test, C_test, Cond_indices_test = test_instance.get_all_items()
this_dataloader_color = test_instance.get_color()
else:
X_train, C_train, Cond_indices_train, _ = load_data(
args.num_batches,
args.batch_size,
args.model_kwargs,
corr=False,
train=True,
mask=mask_bool,
).get_all_items()
test_instance = load_data(
args.num_batches,
args.batch_size,
args.model_kwargs,
corr=False,
train=False,
mask=mask_bool,
)
X_test, C_test, Cond_indices_test = test_instance.get_all_items()
this_dataloader_color = test_instance.get_color()
model = get_model(args.model_fn, args.model_kwargs).to(device)
opt = optim.Adam(model.parameters(), lr=args.lr)
loss_fn = str_to_object(args.loss_fn)
make_plot_encoding_greedy = str_to_object(
"CVAE_testbed.utils.greedy_encoding_plots.make_plot_encoding_greedy")
make_plot_encoding = str_to_object(
"CVAE_testbed.utils.encoding_plots.make_plot_encoding")
make_plot = str_to_object(
"CVAE_testbed.utils.loss_and_image_plots.make_plot")
pca = str_to_object("CVAE_testbed.utils.pca.get_PCA_features")
make_fid_plot = str_to_object("CVAE_testbed.utils.FID_score.make_plot_FID")
if args.data_type == "mnist":
run = str_to_object(
"CVAE_testbed.run_models.run_test_train.run_test_train")
stats = run(
model,
opt,
loss_fn,
device,
args.batch_size,
train_iterator,
test_iterator,
args.n_epochs,
args.model_kwargs,
)
elif args.data_type == "synthetic" or args.data_type == 'aics_features':
run = str_to_object(
"CVAE_testbed.run_models.run_synthetic_test.run_synthetic")
stats, stats_per_dim = run(
args,
X_train,
C_train,
Cond_indices_train,
X_test,
C_test,
Cond_indices_test,
args.n_epochs,
args.loss_fn,
model,
opt,
args.batch_size,
args.gpu_id,
args.model_kwargs,
)
if args.data_type == 'aics_features' or args.data_type == 'synthetic':
print(proj_matrix)
# First load non shuffled data
if proj_matrix is not None:
this_dataloader = load_data(args.num_batches,
args.batch_size,
args.model_kwargs,
shuffle=False,
P=proj_matrix,
train=False)
elif args.data_type == 'aics_features':
this_dataloader = load_data(args.num_batches,
args.batch_size,
args.model_kwargs,
shuffle=False,
train=False)
else:
this_dataloader = load_data(args.num_batches,
args.batch_size,
args.model_kwargs,
shuffle=False,
train=False)
X_non_shuffled, C_non_shuffled, _ = this_dataloader.get_all_items()
# Now check encoding
make_plot_encoding(args, model, stats, X_non_shuffled.clone(),
C_non_shuffled.clone(), this_dataloader_color,
True, proj_matrix)
pca_dataframe = pca(args, this_dataloader, True)
make_plot_encoding_greedy(args, model, stats,
X_non_shuffled.clone(),
C_non_shuffled.clone(), feature_names,
True, proj_matrix)
try:
make_fid_plot(args, model, X_non_shuffled.clone(),
C_non_shuffled.clone())
except:
pass
else:
make_plot_encoding(args, model, stats, X_test, C_test)
this_model = ModelLoader(model, path_save_dir)
this_model.save_model()
path_csv = path_save_dir / Path("costs.csv")
stats.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_csv = path_save_dir / Path("costs_per_dimension.csv")
stats_per_dim.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
make_plot(stats, stats_per_dim, path_save_dir, args)
LOGGER.info(f"Elapsed time: {time.time() - tic:.2f}")
print("saved:", path_save_dir)
def make_plot_FID(args: argparse.Namespace, model, X_test, C_test, save=True):
X_test = X_test.view(-1, X_test.size()[-1])
C_test = C_test.view(-1, C_test.size()[-1])
# X_test = X_test[-1,:]
# C_test = C_test[-1,:]
print(X_test.size(), C_test.size())
sns.set_context("talk")
path_save_dir = Path(args.path_save_dir)
compute_fid = str_to_object(
"CVAE_testbed.run_models.generative_metric.compute_generative_metric_synthetic"
)
try:
this_kwargs = args.model_kwargs["dec_layers"][-1][-1]
except:
this_kwargs = args.model_kwargs["dec_layers"][-1]
# ADD YOUR PATH HERE
csv_greedy_features = pd.read_csv(
'~/Github/cookiecutter/CVAE_testbed/scripts' + args.path_save_dir[1:] +
'/selected_features.csv')
# conds = [i for i in range(this_kwargs)]
conds = [
i for i in csv_greedy_features['selected_feature_number']
if not math.isnan(i)
]
fid_data = {'num_conds': [], 'fid': []}
print(len(conds))
for i in range(len(conds) + 1):
tmp1, tmp2 = torch.split(C_test.clone(),
int(C_test.clone().size()[-1] / 2),
dim=1)
for kk in conds:
tmp1[:, int(kk)], tmp2[:, int(kk)] = 0, 0
cond_d = torch.cat((tmp1, tmp2), 1)
print(len(torch.nonzero(cond_d)))
try:
this_fid = compute_fid(X_test.clone(), cond_d.clone(), args, model,
conds)
except:
this_fid = np.NaN
print('fid', this_fid)
fid_data['num_conds'].append(X_test.size()[-1] - len(conds))
fid_data['fid'].append(this_fid)
try:
conds.pop()
except:
pass
fid_data = pd.DataFrame(fid_data)
fig, ax = plt.subplots(1, 1, figsize=(7 * 4, 5))
sns.lineplot(ax=ax, data=fid_data, x='num_conds', y='fid')
sns.scatterplot(ax=ax,
data=fid_data,
x='num_conds',
y='fid',
s=100,
color=".2")
if save is True:
path_csv = path_save_dir / Path("fid_data.csv")
fid_data.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_save_fig = path_save_dir / Path("fid_score.png")
fig.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
def visualize_encoder_synthetic(args,
model,
conds,
c,
d,
kl_per_lt=None,
kl_vs_rcl=None):
z_means_x, z_means_y = [], []
z_var_x, z_var_y = [], []
model.cuda(args.gpu_id)
with torch.no_grad():
if kl_per_lt is None:
kl_per_lt = {
"latent_dim": [],
"kl_divergence": [],
"num_conds": [],
}
if kl_vs_rcl is None:
kl_vs_rcl = {"num_conds": [], "KLD": [], "RCL": [], "ELBO": []}
all_kl, all_lt = [], []
tmp1, tmp2 = torch.split(d, int(d.size()[-1] / 2), dim=1)
for kk in conds:
tmp1[:, kk], tmp2[:, kk] = 0, 0
cond_d = torch.cat((tmp1, tmp2), 1)
# Run 10 times for resampling
my_recon_list, my_z_means_list, my_log_var_list = [], [], []
for resample in range(10):
recon_batch, z_means, log_var = model(
c.clone().cuda(args.gpu_id),
cond_d.clone().cuda(args.gpu_id))
my_recon_list.append(recon_batch)
my_z_means_list.append(z_means)
my_log_var_list.append(log_var)
recon_batch = torch.mean(torch.stack(my_recon_list), dim=0)
z_means = torch.mean(torch.stack(my_z_means_list), dim=0)
log_var = torch.mean(torch.stack(my_log_var_list), dim=0)
loss_fn = str_to_object(args.loss_fn)
elbo_loss_total, rcl_per_lt_temp_total, kl_per_lt_temp_total, _, _ = loss_fn(
c.cuda(args.gpu_id), recon_batch.cuda(args.gpu_id), z_means,
log_var, args)
# print('conds is', conds, elbo_loss_total, rcl_per_lt_temp_total, kl_per_lt_temp_total)
kl_vs_rcl['num_conds'].append(c.size()[-1] - len(conds))
kl_vs_rcl['KLD'].append(kl_per_lt_temp_total.item())
kl_vs_rcl['RCL'].append(rcl_per_lt_temp_total.item())
kl_vs_rcl['ELBO'].append(elbo_loss_total.item())
for ii in range(z_means.size()[-1]):
elbo_loss, rcl_per_lt_temp, kl_per_lt_temp, _, _ = loss_fn(
c.cuda(args.gpu_id), recon_batch.cuda(args.gpu_id),
z_means[:, ii], log_var[:, ii], args)
# print(elbo_loss.item(), rcl_per_lt_temp.item(), kl_per_lt_temp.item())
all_kl = np.append(all_kl, kl_per_lt_temp.item())
all_lt.append(ii)
kl_per_lt["num_conds"].append(c.size()[-1] - len(conds))
kl_per_lt["latent_dim"].append(ii)
kl_per_lt["kl_divergence"].append(kl_per_lt_temp.item())
all_kl, all_lt = list(zip(*sorted(zip(all_kl, all_lt))))
all_kl = list(all_kl)
all_lt = list(all_lt)
z_means_x = np.append(z_means_x,
z_means[:, all_lt[-1]].data.cpu().numpy())
z_means_y = np.append(z_means_y,
z_means[:, all_lt[-2]].data.cpu().numpy())
z_var_x = np.append(z_var_x, log_var[:, all_lt[-1]].data.cpu().numpy())
z_var_y = np.append(z_var_y, log_var[:, all_lt[-2]].data.cpu().numpy())
return z_means_x, z_means_y, kl_per_lt, z_var_x, z_var_y, kl_vs_rcl
def train(
args,
epoch,
loss_fn,
X_train,
C_train,
Cond_indices_train,
batch_size,
model,
optimizer,
gpu_id,
model_kwargs,
):
model.train()
train_loss, rcl_loss, kld_loss = 0, 0, 0
rcl_per_condition_loss, kld_per_condition_loss = (
torch.zeros(X_train.size()[-1] + 1),
torch.zeros(X_train.size()[-1] + 1),
)
batch_rcl, batch_kld = (
torch.empty([0]),
torch.empty([0]),
)
batch_length = 0
for j, i in enumerate(X_train):
optimizer.zero_grad()
c, d, cond_labels = X_train[j], C_train[j], Cond_indices_train[j]
if args.model_fn == 'CVAE_testbed.models.CVAE_baseline_2.CVAE':
recon_batch, z1, mu, log_var, mu2, log_var2, z2, z1_prior = model(
c.cuda(gpu_id), d.cuda(gpu_id))
loss_fn = str_to_object(args.loss_fn)
loss, rcl, kld, rcl_per_element, kld_per_element = loss_fn(
c.cuda(gpu_id), recon_batch.cuda(gpu_id), z1, z1_prior, z2,
[mu, log_var], [mu2, log_var2], args)
else:
recon_batch, mu, log_var = model(c.cuda(gpu_id), d.cuda(gpu_id))
loss_fn = str_to_object(args.loss_fn)
loss, rcl, kld, rcl_per_element, kld_per_element = loss_fn(
c.cuda(gpu_id), recon_batch.cuda(gpu_id), mu, log_var, args)
# print(loss, rcl, kld)
loss.backward()
train_loss += loss.item()
rcl_loss += rcl.item()
kld_loss += kld.item()
for jj, ii in enumerate(torch.unique(cond_labels)):
this_cond_positions = cond_labels == ii
batch_rcl = torch.cat(
[
batch_rcl.cuda(gpu_id),
torch.sum(rcl_per_element[this_cond_positions],
dim=0).view(1, -1),
],
0,
)
batch_kld = torch.cat(
[
batch_kld.cuda(gpu_id),
torch.sum(kld_per_element[this_cond_positions],
dim=0).view(1, -1),
],
0,
)
batch_length += 1
this_cond_rcl = torch.sum(rcl_per_element[this_cond_positions])
this_cond_kld = torch.sum(kld_per_element[this_cond_positions])
rcl_per_condition_loss[jj] += this_cond_rcl.item()
kld_per_condition_loss[jj] += this_cond_kld.item()
optimizer.step()
num_batches = len(X_train)
# print(num_batches)
print("====> Epoch: {} Train loss: {:.4f}".format(epoch, train_loss /
num_batches))
print("====> Train RCL loss: {:.4f}".format(rcl_loss / num_batches))
print("====> Train KLD loss: {:.4f}".format(kld_loss / num_batches))
batch_rcl, batch_kld = batch_rcl, batch_kld
return (
train_loss / num_batches,
rcl_loss / num_batches,
kld_loss / num_batches,
rcl_per_condition_loss / num_batches,
kld_per_condition_loss / num_batches,
batch_rcl,
batch_kld,
batch_length,
)
def make_plot_encoding(
args: argparse.Namespace,
model,
df: pd.DataFrame,
c,
d,
this_dataloader_color=None,
save=True,
proj_matrix=None
) -> None:
"""
c and d are X_test and C_test
"""
sns.set_context("talk")
path_save_dir = Path(args.path_save_dir)
vis_enc = str_to_object(
"CVAE_testbed.metrics.visualize_encoder.visualize_encoder_synthetic"
)
try:
conds = [i for i in range(args.model_kwargs["dec_layers"][-1][-1])]
except:
conds = [i for i in range(args.model_kwargs["dec_layers"][-1])]
try:
latent_dims = args.model_kwargs["vae_layers"][-1][-1]
except:
latent_dims = args.model_kwargs["enc_layers"][-1]
fig, ax = plt.subplots(1, 1, figsize=(7, 5))
sns.lineplot(ax=ax, data=df, x="epoch", y="total_train_ELBO")
sns.lineplot(ax=ax, data=df, x="epoch", y="total_test_ELBO")
ax.set_ylim([0, df.total_test_ELBO.quantile(0.95)])
ax.legend(["Train loss", "Test loss"])
ax.set_ylabel('Loss')
ax.set_title("Actual ELBO (no beta) vs epoch")
if save is True:
path_save_fig = path_save_dir / Path("ELBO.png")
fig.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
fig, (ax1, ax, ax2, ax3) = plt.subplots(1, 4, figsize=(7 * 4, 5))
fig2 = plt.figure(figsize=(12, 10))
bax = brokenaxes(xlims=((0, latent_dims-50), (latent_dims - 4, latent_dims)), hspace=0.15)
if "total_train_losses" in df.columns:
sns.lineplot(ax=ax1, data=df, x="epoch", y="total_train_losses")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss")
if "total_test_losses" in df.columns:
sns.lineplot(ax=ax1, data=df, x="epoch", y="total_test_losses")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss")
ax1.set_ylim([0, df.total_test_losses.quantile(0.95)])
ax1.legend(["Train loss", "Test loss"])
ax1.set_title("ELBO (beta*KLD + RCL) vs epoch")
try:
this_kwargs = args.model_kwargs["dec_layers"][-1][-1]
except:
this_kwargs = args.model_kwargs["dec_layers"][-1]
conds = [i for i in range(this_kwargs)]
# if len(conds) > 20:
# conds = [i for i in conds if i%20 == 0]
if args.post_plot_kwargs["latent_space_colorbar"] == "yes":
# color = this_dataloader.get_color()
color = this_dataloader_color
else:
color = None
for i in range(len(conds) + 1):
# print('inside main plot encoding', i, len(conds) + 1)
if i == 0:
z_means_x, z_means_y, kl_per_lt, _, _, kl_vs_rcl = vis_enc(
args,
model,
conds,
c[-1, :].clone(),
d[-1, :].clone(),
kl_per_lt=None,
kl_vs_rcl=None
)
ax.scatter(
z_means_x,
z_means_y,
marker=".",
s=30,
label=str(i)
)
if color is not None:
colormap_plot(
path_save_dir,
c[-1, :].clone(), z_means_x,
z_means_y, color,
conds)
else:
z_means_x, z_means_y, kl_per_lt, _, _, kl_vs_rcl = vis_enc(
args,
model,
conds,
c[-1, :].clone(),
d[-1, :].clone(),
kl_per_lt,
kl_vs_rcl
)
ax.scatter(
z_means_x,
z_means_y,
marker=".",
s=30,
label=str(i)
)
if color is not None:
colormap_plot(
path_save_dir,
c[-1, :].clone(),
z_means_x,
z_means_y,
color,
conds
)
try:
conds.pop()
except:
pass
kl_per_lt = pd.DataFrame(kl_per_lt)
kl_vs_rcl = pd.DataFrame(kl_vs_rcl)
if save is True:
path_csv = path_save_dir / Path("encoding_kl_per_lt.csv")
kl_per_lt.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
path_csv = path_save_dir / Path("encoding_kl_vs_rcl.csv")
kl_vs_rcl.to_csv(path_csv)
LOGGER.info(f"Saved: {path_csv}")
ax.set_title("Latent space")
ax.legend()
conds = [i for i in range(this_kwargs)]
# if len(conds) > 20:
# conds = [i for i in conds if i%20 == 0]
for i in range(len(conds) + 1):
tmp = kl_per_lt.loc[kl_per_lt["num_conds"] == c.size()[-1] - len(conds)]
tmp_2 = kl_vs_rcl.loc[kl_vs_rcl["num_conds"] == c.size()[-1] - len(conds)]
tmp = tmp.sort_values(
by="kl_divergence",
ascending=False
)
tmp = tmp.reset_index(drop=True)
x = tmp.index.values
y = tmp.iloc[:, 1].values
sns.lineplot(
ax=ax2,
data=tmp,
x=tmp.index,
y="kl_divergence",
label=str(i),
legend='brief'
)
bax.plot(x, y)
ax3.scatter(tmp_2['RCL'].mean(), tmp_2['KLD'].mean(), label=str(i))
# sns.scatterplot(
# ax=ax3,
# data=tmp,
# x="rcl",
# y="kl_divergence",
# label=str(i),
# legend='brief'
# )
try:
conds.pop()
except:
pass
ax2.set_xlabel("Latent dimension")
ax2.set_ylabel("KLD")
ax2.set_title("KLD per latent dim")
ax3.set_xlabel("MSE")
ax3.set_ylabel("KLD")
ax3.set_title("MSE vs KLD")
# bax.legend(loc="best")
bax.set_xlabel("Latent dimension")
bax.set_ylabel("KLD")
bax.set_title("KLD per latent dim")
conds = [i for i in range(this_kwargs)]
if len(conds) > 30:
ax.get_legend().remove()
ax2.get_legend().remove()
if save is True:
path_save_fig = path_save_dir / Path("encoding_test_plots.png")
fig.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")
if save is True:
path_save_fig = path_save_dir / Path("brokenaxes_KLD_per_dim.png")
fig2.savefig(path_save_fig, bbox_inches="tight")
LOGGER.info(f"Saved: {path_save_fig}")