"""Evaluates the model"""
import argparse
import logging
import os
import numpy as np
import pandas as pd
import torch
from torch.autograd import Variable
import model.data_loader as data_loader
import model.net as net
import utils
from sklearn import linear_model
def evaluate(model, loss_fn, dataloader, metrics, params, setting, epoch, writer=None):
"""Evaluate the model on `num_steps` batches.
Args:
model: (torch.nn.Module) the neural network
loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch
dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches data
metrics: (dict) a dictionary of functions that compute a metric using the output and labels of each batch
params: (Params) hyperparameters
num_steps: (int) number of batches to train on, each of size params.batch_size
covar_mode: (bool) include covariate data in dataloader
"""
# set model to evaluation mode
model.eval()
model.to(params.device)
# summary for current eval loop
summ = []
preds = [] # for saving last predictions
bn_activations = []
# create storate for tensors for OLS after minibatches
Xhats = []
Zhats = []
# for counterfactuals
if setting.counterfactuals:
y0_hats = []
y1_hats = []
# compute metrics over the dataset
for batch in dataloader:
summary_batch = {}
batch = {k: v.to(params.device) for k, v in batch.items()}
img_batch = batch["image"].to(params.device, non_blocking=True)
labels_batch = batch["label"].to(params.device, non_blocking=True)
if setting.covar_mode and epoch > params.suppress_t_epochs:
data_batch = batch["t"].to(params.device, non_blocking=True).view(-1,1)
else:
data_batch = torch.zeros((params.batch_size, 1), requires_grad=False).to(params.device, non_blocking=True)
if params.multi_task:
# x_target_batch = Variable(batch["x"].to(params.device)).type(torch.cuda.LongTensor)
x_target_batch = batch["x"].to(params.device)
y_target_batch = batch["y"].to(params.device)
labels_batch = {'x': x_target_batch, 'y': y_target_batch}
# compute model output
# output_batch, bn_batch = model(img_batch, data_batch)
output_batch = model(img_batch, data_batch, epoch)
# calculate loss
if setting.fase == "feature":
# calculate loss for z directly, to get clear how well this can be measured
loss_fn_z = torch.nn.MSELoss()
loss_z = loss_fn_z(output_batch["y"].squeeze(), batch["z"])
loss = loss_z
summary_batch["loss_z"] = loss_z.item()
else:
loss_fn_y = torch.nn.MSELoss()
loss_y = loss_fn_y(output_batch["y"].squeeze(), batch["y"])
loss = loss_y
summary_batch["loss_y"] = loss_y.item()
# calculate loss for colllider x
loss_fn_x = torch.nn.MSELoss()
loss_x = loss_fn_x(output_batch["bnx"].squeeze(), batch["x"])
summary_batch["loss_x"] = loss_x.item()
if not params.alpha == 1:
# possibly weigh down contribution of estimating x
loss_x *= params.alpha
summary_batch["loss_x_weighted"] = loss_x.item()
# add x loss to total loss
loss += loss_x
# add least squares regression on final layer
if params.do_least_squares:
X = batch["x"].view(-1,1)
t = batch["t"].view(-1,1)
Z = output_batch["bnz"]
if Z.ndimension() == 1:
Z.unsqueeze_(1)
Xhat = output_batch["bnx"]
# add intercept
Zi = torch.cat([torch.ones_like(t), Z], 1)
# add treatment info
Zt = torch.cat([Zi, t], 1)
Y = batch["y"].view(-1,1)
# regress y on final layer, without x
betas_y = net.cholesky_least_squares(Zt, Y, intercept=False)
y_hat = Zt.matmul(betas_y).view(-1,1)
mse_y = ((Y - y_hat)**2).mean()
summary_batch["regr_b_t"] = betas_y[-1].item()
summary_batch["regr_loss_y"] = mse_y.item()
# regress x on final layer without x
betas_x = net.cholesky_least_squares(Zi, Xhat, intercept=False)
x_hat = Zi.matmul(betas_x).view(-1,1)
mse_x = ((Xhat - x_hat)**2).mean()
# store all tensors for single pass after epoch
Xhats.append(Xhat.detach().cpu())
Zhats.append(Z.detach().cpu())
summary_batch["regr_loss_x"] = mse_x.item()
# add loss_bn only after n epochs
if params.bottleneck_loss and epoch > params.bn_loss_lag_epochs:
# only add to loss when bigger than margin
if params.bn_loss_margin_type == "dynamic-mean":
# for each batch, calculate loss of just using mean for predicting x
mse_x_mean = ((X - X.mean())**2).mean()
loss_bn = torch.max(torch.zeros_like(mse_x), mse_x_mean - mse_x)
elif params.bn_loss_margin_type == "fixed":
mse_diff = params.bn_loss_margin - mse_x
loss_bn = torch.max(torch.zeros_like(mse_x), mse_diff)
else:
raise NotImplementedError(f'bottleneck loss margin type not implemented: {params.bn_loss_margin_type}')
# possibly reweigh bottleneck loss and add to total loss
summary_batch["loss_bn"] = loss_bn.item()
# note is this double?
if loss_bn > params.bn_loss_margin:
loss_bn *= params.bottleneck_loss_wt
loss += loss_bn
# generate counterfactual predictions
if setting.counterfactuals:
batch_t0 = Variable(torch.zeros_like(data_batch).to(torch.float32), requires_grad=False).to(params.device)
batch_t1 = Variable(torch.ones_like(data_batch).to(torch.float32), requires_grad=False).to(params.device)
y0_batch = model(img_batch, batch_t0)
y1_batch = model(img_batch, batch_t1)
y0_hats.append(y0_batch["y"].detach().cpu().numpy())
y1_hats.append(y1_batch["y"].detach().cpu().numpy())
# write out activations of bottleneck layer
if params.multi_task:
bn_activations.append(output_batch["bnz"])
else:
bn_activations.append(output_batch["bn"])
# extract data from torch Variable, move to cpu, convert to numpy arrays
if (len(setting.outcome) > 1) or params.multi_task:
for var, batch in labels_batch.items():
labels_batch[var] = batch.data.cpu().numpy()
else:
labels_batch = labels_batch.data.cpu().numpy()
# compute all metrics on this batch
data_batch = data_batch.data.cpu().numpy()
for var, batch in output_batch.items():
output_batch[var] = batch.detach().cpu().numpy()
if params.multi_task:
metrics_xy = {m: net.all_metrics[m] for m in setting.metrics_xy}
for var, batch in labels_batch.items():
for metric, metric_fn in metrics_xy.items():
summary_batch[metric+"_"+var] = metric_fn(setting, model, output_batch[var], labels_batch[var], data_batch)
if "b_t" in setting.metrics:
summary_batch["b_t"] = net.all_metrics["b_t"](setting, model, None, None)
else:
NotImplementedError
# summary_batch = {metric: metrics[metric](setting, model, output_batch[setting.outcome[0]], labels_batch, data_batch)
# for metric in metrics}
summary_batch["loss"] = loss.item()
summ.append(summary_batch)
#if "y" in setting.outcome:
preds.append(output_batch["y"])
#else:
# preds.append(output_batch[setting.outcome[0]])
# compute mean of all metrics in summary
metrics_mean = {metric:np.nanmean([x[metric] for x in summ]) for metric in summ[0]}
# if "ate" in setting.metrics:
# metrics_mean["ate"] = all_metrics["ate"](setting, model, preds, )
if params.save_bn_activations:
# write out batch activations
bn_activations = torch.cat(bn_activations, 0).detach().cpu().numpy()
writer.add_histogram("bn_activations", bn_activations, epoch+1)
# get means and covariances
if "bottleneck_loss" in setting.metrics:
bn_means = bn_activations.mean(dim=0)
bn_sds = bn_activations.std(dim=0)
bn_cov = net.cov(bn_activations)
bn_offdiags = net.get_of_diag(bn_cov.detach().cpu().numpy())
writer.add_histogram("bn_covariances", bn_offdiags, epoch+1)
# export predictions
preds = np.vstack([x.reshape(-1,1) for x in preds])
writer.add_histogram('predictions', preds, epoch+1)
labels = dataloader.dataset.df[setting.outcome[0]].values.astype(np.float32)
# predict individual treatment effects (only worth-while when there is an interaction with t)
if setting.counterfactuals:
y0_hats = np.vstack(y0_hats)
y1_hats = np.vstack(y1_hats)
ite_hats = y1_hats - y0_hats
metrics_mean["ite_mean"] = ite_hats.mean()
y0s = dataloader.dataset.df["y0"].values.astype(np.float32)
y1s = dataloader.dataset.df["y1"].values.astype(np.float32)
ites = y1s - y0s
metrics_mean["pehe"] = np.sqrt(np.mean(np.power((ite_hats - ites), 2)))
metrics_mean["loss_y1"] = ((y1s - y1_hats)**2).mean()
metrics_mean["loss_y0"] = ((y0s - y0_hats)**2).mean()
# in case of single last layer where x is part of, do regression on this layer
if params.bn_place == "single-regressor" and params.do_least_squares:
Xhat = torch.cat(Xhats, 0).view(-1,1).float()
Zhat = torch.cat(Zhats, 0).float()
t = torch.tensor(dataloader.dataset.df["t"].values).view(-1,1).float()
Y = torch.tensor(dataloader.dataset.df["y"].values).view(-1,1).float()
betas_bias = model.betas_bias.cpu()
betas_causal = model.betas_causal.cpu()
y_hat_bias = torch.cat([torch.ones_like(t), Xhat, Zhat, t], 1).matmul(betas_bias).view(-1,1)
y_hat_causal = torch.cat([torch.ones_like(t), Zhat, t], 1).matmul(betas_causal).view(-1,1)
reg_mse_bias = ((y_hat_bias - Y)**2).mean()
reg_mse_causal = ((y_hat_causal - Y)**2).mean()
metrics_mean["regr_bias_loss_y"] = reg_mse_bias
metrics_mean["regr_causal_loss_y"] = reg_mse_causal
if setting.counterfactuals:
y0_hat_bias = torch.cat([torch.ones_like(t), Xhat, Zhat, torch.zeros_like(t)], 1).matmul(betas_bias).view(-1,1)
y1_hat_bias = torch.cat([torch.ones_like(t), Xhat, Zhat, torch.ones_like(t)], 1).matmul(betas_bias).view(-1,1)
y0_hat_causal = torch.cat([torch.ones_like(t), Zhat, torch.zeros_like(t)], 1).matmul(betas_causal).view(-1,1)
y1_hat_causal = torch.cat([torch.ones_like(t), Zhat, torch.ones_like(t)], 1).matmul(betas_causal).view(-1,1)
ite_hats_bias = y1_hat_bias - y0_hat_bias
ite_hats_causal = y1_hat_causal - y0_hat_causal
writer.add_scalars("pehe", {"regr_bias": np.sqrt(((ite_hat_bias - ites)**2).mean())}, epoch+1)
writer.add_scalars("pehe", {"regr_causal": np.sqrt(((ite_hat_causal - ites)**2).mean())}, epoch+1)
writer.add_scalars("loss_y1", {"regr_bias": ((y1s - y1_hat_bias)**2).mean()}, epoch+1)
writer.add_scalars("loss_y0", {"regr_bias": ((y0s - y0_hat_bias)**2).mean()}, epoch+1)
writer.add_scalars("loss_y1", {"regr_causal": ((y1s - y1_hat_causal)**2).mean()}, epoch+1)
writer.add_scalars("loss_y0", {"regr_causal": ((y0s - y0_hat_causal)**2).mean()}, epoch+1)
outtensors = {
'bn_activations': bn_activations,
'predictions': preds,
'xhat': np.vstack(Xhats)
}
metrics_string = " ; ".join("{}: {:05.3f}".format(k, v) for k, v in metrics_mean.items())
logging.info("- Eval metrics : " + metrics_string)
return metrics_mean, outtensors
Datasets
Models
