"""Train the model"""
import argparse
import logging
import os, shutil
import numpy as np
import pandas as pd
from sklearn.utils.class_weight import compute_class_weight
import torch
import torch.optim as optim
import torchvision.models as models
from torch.autograd import Variable
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
# from torchsummary import summary
import utils
import json
import model.net as net
import model.data_loader as data_loader
from evaluate import evaluate
parser = argparse.ArgumentParser()
parser.add_argument('--data-dir', default='data', help="Directory containing the dataset")
parser.add_argument('--model-dir', default='experiments', help="Directory containing params.json")
parser.add_argument('--setting-dir', default='settings', help="Directory with different settings")
parser.add_argument('--setting', default='collider-prognosticfactor', help="Directory contain setting.json, experimental setting, data-generation, regression model etc")
parser.add_argument('--fase', default='xybn', help='fase of training model, see manuscript for details. x, y, xy, bn, or feature')
parser.add_argument('--experiment', default='', help="Manual name for experiment for logging, will be subdir of setting")
parser.add_argument('--restore-file', default=None,
help="Optional, name of the file in --model_dir containing weights to reload before \
training") # 'best' or 'train'
parser.add_argument('--restore-last', action='store_true', help="continue a last run")
parser.add_argument('--restore-warm', action='store_true', help="continue on the run called 'warm-start.pth'")
parser.add_argument('--use-last', action="store_true", help="use last state dict instead of 'best' (use for early stopping manually)")
parser.add_argument('--cold-start', action='store_true', help="ignore previous state dicts (weights), even if they exist")
parser.add_argument('--warm-start', dest='cold_start', action='store_false', help="start from previous state dict")
parser.add_argument('--disable-cuda', action='store_true', help="Disable Cuda")
parser.add_argument('--no-parallel', action="store_false", help="no multiple GPU", dest="parallel")
parser.add_argument('--parallel', action="store_true", help="multiple GPU", dest="parallel")
parser.add_argument('--gpu', default=0, type=int, help='if not running in parallel (=all gpus), only use this gpu')
parser.add_argument('--intercept', action="store_true", help="dummy run for getting intercept baseline results")
# parser.add_argument('--visdom', action='store_true', help='generate plots with visdom')
# parser.add_argument('--novisdom', dest='visdom', action='store_false', help='dont plot with visdom')
parser.add_argument('--monitor-grads', action='store_true', help='keep track of mean norm of gradients')
parser.set_defaults(parallel=False, cold_start=True, use_last=False, intercept=False, restore_last=False, save_preds=False,
monitor_grads=False, restore_warm=False
# visdom=False
)
def train(model, optimizer, loss_fn, dataloader, metrics, params, setting, writer=None, epoch=None):
"""Train the model on `num_steps` batches
Args:
model: (torch.nn.Module) the neural network
optimizer: (torch.optim) optimizer for parameters of model
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 training 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
"""
global train_tensor_keys, logdir
# set model to training mode
model.train()
# summary for current training loop and a running average object for loss
summ = []
loss_avg = utils.RunningAverage()
# create storate for tensors for OLS after minibatches
ts = []
Xs = []
Xtrues = []
Ys = []
Xhats = []
Yhats = []
Zhats = []
# Use tqdm for progress bar
with tqdm(total=len(dataloader)) as progress_bar:
for i, batch in enumerate(dataloader):
summary_batch = {}
# put batch on cuda
batch = {k: v.to(params.device) for k, v in batch.items()}
if not (setting.covar_mode and epoch > params.suppress_t_epochs):
batch["t"] = torch.zeros_like(batch['t'])
Xs.append(batch['x'].detach().cpu())
Xtrues.append(batch['x_true'].detach().cpu())
# compute model output and loss
output_batch = model(batch['image'], batch['t'].view(-1,1), epoch)
Yhats.append(output_batch['y'].detach().cpu())
# calculate loss
if args.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())
ts.append(t.detach().cpu())
Ys.append(Y.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
# perform parameter update
optimizer.zero_grad()
loss.backward()
optimizer.step()
summary_batch['loss'] = loss.item()
summ.append(summary_batch)
# if necessary, write out tensors
if params.monitor_train_tensors and (epoch % params.save_summary_steps == 0):
tensors = {}
for tensor_key in train_tensor_keys:
if tensor_key in batch.keys():
tensors[tensor_key] = batch[tensor_key].squeeze().numpy()
elif tensor_key.endswith("hat"):
tensor_key = tensor_key.split("_")[0]
if tensor_key in output_batch.keys():
tensors[tensor_key+"_hat"] = output_batch[tensor_key].detach().cpu().squeeze().numpy()
else:
assert False, f"key not found: {tensor_key}"
# print(tensors)
df = pd.DataFrame.from_dict(tensors, orient='columns')
df["epoch"] = epoch
with open(os.path.join(logdir, 'train-tensors.csv'), 'a') as f:
df[["epoch"]+train_tensor_keys].to_csv(f, header=False)
# update the average loss
loss_avg.update(loss.item())
progress_bar.set_postfix(loss='{:05.3f}'.format(loss_avg()))
progress_bar.update()
# visualize gradients
if epoch % params.save_summary_steps == 0 and args.monitor_grads:
abs_gradients = {}
for name, param in model.named_parameters():
try: # patch here, there were names / params that were 'none'
abs_gradients[name] = np.abs(param.grad.cpu().numpy()).mean()
writer.add_histogram("grad-"+name, param.grad, epoch)
writer.add_scalars("mean-abs-gradients", abs_gradients, epoch)
except:
pass
# compute mean of all metrics in summary
metrics_mean = {metric:np.nanmean([x[metric] for x in summ]) for metric in summ[0]}
# collect tensors
Xhat = torch.cat(Xhats,0).view(-1,1)
Yhat = torch.cat(Yhats,0).view(-1,1)
Zhat = torch.cat(Zhats,0)
t = torch.cat(ts,0)
X = torch.cat(Xs,0)
Xtrue= torch.cat(Xtrues,0)
Y = torch.cat(Ys,0)
if params.do_least_squares:
# after the minibatches, do a single OLS on the whole data
Zi = torch.cat([torch.ones_like(t), Zhat], 1)
# add treatment info
Zt = torch.cat([Zi, t], 1)
# add x for biased version
XZt = torch.cat([torch.ones_like(t), Xhat, Zhat, t], 1)
betas_y_bias = net.cholesky_least_squares(XZt, Y, intercept=False)
betas_y_causal = net.cholesky_least_squares(Zt, Y, intercept=False)
model.betas_bias = betas_y_bias
model.betas_causal = betas_y_causal
metrics_mean["regr_bias_coef_t"] = betas_y_bias.squeeze()[-1]
metrics_mean["regr_bias_coef_z"] = betas_y_bias.squeeze()[-2]
metrics_mean["regr_causal_coef_t"] = betas_y_causal.squeeze()[-1]
metrics_mean["regr_causal_coef_z"] = betas_y_causal.squeeze()[-2]
# create some plots
xx_scatter = net.make_scatter_plot(X.numpy(), Xhat.numpy(), xlabel='x', ylabel='xhat')
xtruex_scatter= net.make_scatter_plot(Xtrue.numpy(), Xhat.numpy(), xlabel='xtrue', ylabel='xhat')
xyhat_scatter = net.make_scatter_plot(X.numpy(), Yhat.numpy(), c=t.numpy(), xlabel='x', ylabel='yhat')
yy_scatter = net.make_scatter_plot(Y.numpy(), Yhat.numpy(), c=t.numpy(), xlabel='y', ylabel='yhat')
writer.add_figure('x-xhat/train', xx_scatter, epoch+1)
writer.add_figure('xtrue-xhat/train', xtruex_scatter, epoch+1)
writer.add_figure('x-yhat/train', xyhat_scatter, epoch+1)
writer.add_figure('y-yhat/train', yy_scatter, epoch+1)
metrics_string = " ; ".join("{}: {:05.3f}".format(k, v) for k, v in metrics_mean.items())
logging.info("- Train metrics: " + metrics_string)
return metrics_mean
def train_and_evaluate(model, train_dataloader, val_dataloader, optimizer, loss_fn, metrics, params, setting, args,
writer=None, logdir=None, restore_file=None):
"""Train the model and evaluate every epoch.
Args:
model: (torch.nn.Module) the neural network
train_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches training data
val_dataloader: (DataLoader) a torch.utils.data.DataLoader object that fetches validation data
optimizer: (torch.optim) optimizer for parameters of model
loss_fn: a function that takes batch_output and batch_labels and computes the loss for the batch
metrics: (dict) a dictionary of functions that compute a metric using mnisthe output and labels of each batch
params: (Params) hyperparameters
model_dir: (string) directory containing config, weights and log
restore_file: (string) optional- name of file to restore from (withoutmnistits extension .pth.tar)
covar_mode: (bool) does the data-loader give back covariates / additional data
"""
# setup directories for data
setting_home = setting.home
if not args.fase == "feature":
data_dir = os.path.join(setting_home, "data")
else:
if setting.mode3d:
data_dir = "data"
else:
data_dir = "slices"
covar_mode = setting.covar_mode
x_frozen = False
best_val_metric = 0.0
if "loss" in setting.metrics[0]:
best_val_metric = 1.0e6
val_preds = np.zeros((len(val_dataloader.dataset), params.num_epochs))
for epoch in range(params.num_epochs):
# Run one epoch
logging.info(f"Epoch {epoch+1}/{params.num_epochs}; setting: {args.setting}, fase {args.fase}, experiment: {args.experiment}")
# compute number of batches in one epoch (one full pass over the training set)
train_metrics = train(model, optimizer, loss_fn, train_dataloader, metrics, params, setting, writer, epoch)
print(train_metrics)
for metric_name in train_metrics.keys():
metric_vals = {'train': train_metrics[metric_name]}
writer.add_scalars(metric_name, metric_vals, epoch+1)
# for name, param in model.named_parameters():
# writer.add_histogram(name, param.clone().cpu().data.numpy(), epoch+1)
if epoch % params.save_summary_steps == 0:
# Evaluate for one epoch on validation set
valid_metrics, outtensors = evaluate(model, loss_fn, val_dataloader, metrics, params, setting, epoch, writer)
valid_metrics["intercept"] = model.regressor.fc.bias.detach().cpu().numpy()
print(valid_metrics)
for name, module in model.regressor.named_children():
if name == "t":
valid_metrics["b_t"] = module.weight.detach().cpu().numpy()
elif name == "zt":
weights = module.weight.detach().cpu().squeeze().numpy().reshape(-1)
for i, weight in enumerate(weights):
valid_metrics["b_zt"+str(i)] = weight
else:
pass
for metric_name in valid_metrics.keys():
metric_vals = {'valid': valid_metrics[metric_name]}
writer.add_scalars(metric_name, metric_vals, epoch+1)
# create plots
val_df = val_dataloader.dataset.df
xx_scatter = net.make_scatter_plot(val_df.x.values, outtensors['xhat'], xlabel='x', ylabel='xhat')
xtruex_scatter= net.make_scatter_plot(val_df.x_true.values, outtensors['xhat'], xlabel='x', ylabel='xhat')
xyhat_scatter = net.make_scatter_plot(val_df.x.values, outtensors['predictions'], c=val_df.t, xlabel='x', ylabel='yhat')
zyhat_scatter = net.make_scatter_plot(val_df.z.values, outtensors['predictions'], c=val_df.t, xlabel='z', ylabel='yhat')
yy_scatter = net.make_scatter_plot(val_df.y.values, outtensors['predictions'], c=val_df.t, xlabel='yhat', ylabel='y')
writer.add_figure('x-xhat/valid', xx_scatter, epoch+1)
writer.add_figure('xtrue-xhat/valid', xtruex_scatter, epoch+1)
writer.add_figure('x-yhat/valid', xyhat_scatter, epoch+1)
writer.add_figure('z-yhat/valid', zyhat_scatter, epoch+1)
writer.add_figure('y-yhat/valid', yy_scatter, epoch+1)
if params.save_preds:
# writer.add_histogram("predictions", preds)
if setting.num_classes == 1:
val_preds[:, epoch] = np.squeeze(outtensors['predictions'])
# write preds to file
pred_fname = os.path.join(setting.home, setting.fase+"-fase", "preds_val.csv")
with open(pred_fname, 'ab') as f:
np.savetxt(f, preds.T, newline="")
np.save(os.path.join(setting.home, setting.fase+"-fase", "preds.npy"), preds)
else:
val_metric = valid_metrics[setting.metrics[0]]
if "loss" in str(setting.metrics[0]):
is_best = val_metric<=best_val_metric
else:
is_best = val_metric>=best_val_metric
# Save weights
state_dict = model.state_dict()
optim_dict = optimizer.state_dict()
state = {
'epoch': epoch+1,
'state_dict': state_dict,
'optim_dict': optim_dict
}
utils.save_checkpoint(state,
is_best=is_best,
checkpoint=logdir)
# If best_eval, best_save_path
valid_metrics["epoch"] = epoch
if is_best:
logging.info("- Found new best {}: {:.3f}".format(setting.metrics[0], val_metric))
best_val_metric = val_metric
# Save best val metrics in a json file in the model directory
best_json_path = os.path.join(logdir, "metrics_val_best_weights.json")
utils.save_dict_to_json(valid_metrics, best_json_path)
# Save latest val metrics in a json file in the model directory
last_json_path = os.path.join(logdir, "metrics_val_last_weights.json")
utils.save_dict_to_json(valid_metrics, last_json_path)
# final evaluation
writer.export_scalars_to_json(os.path.join(logdir, "all_scalars.json"))
if args.save_preds:
np.save(os.path.join(setting.home, setting.fase + "-fase", "val_preds.npy"), val_preds)
if __name__ == '__main__':
# Load the parameters from json file
args = parser.parse_args()
# Load information from last setting if none provided:
last_defaults = utils.Params("last-defaults.json")
if args.setting == "":
print("using last default setting")
args.setting = last_defaults.dict["setting"]
for param, value in last_defaults.dict.items():
print("{}: {}".format(param, value))
else:
with open("last-defaults.json", "r+") as jsonFile:
defaults = json.load(jsonFile)
tmp = defaults["setting"]
defaults["setting"] = args.setting
jsonFile.seek(0) # rewind
json.dump(defaults, jsonFile)
jsonFile.truncate()
# setup visdom environment
# if args.visdom:
# from visdom import Visdom
# viz = Visdom(env=f"lidcr_{args.setting}_{args.fase}_{args.experiment}")
# load setting (data generation, regression model etc)
setting_home = os.path.join(args.setting_dir, args.setting)
setting = utils.Params(os.path.join(setting_home, "setting.json"))
setting.home = setting_home
# when not specified in call, grab model specification from setting file
if setting.cnn_model == "":
json_path = os.path.join(args.model_dir, "t-suppression", args.experiment+".json")
else:
json_path = os.path.join(args.model_dir, setting.cnn_model, 'params.json')
assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)
if not os.path.exists(os.path.join(setting.home, args.fase + "-fase")):
os.makedirs(os.path.join(setting.home, args.fase + "-fase"))
shutil.copy(json_path, os.path.join(setting_home, args.fase + "-fase", "params.json"))
params = utils.Params(json_path)
# covar_mode = setting.covar_mode
# mode3d = setting.mode3d
parallel = args.parallel
params.device = None
if not args.disable_cuda and torch.cuda.is_available():
params.device = torch.device('cuda')
params.cuda = True
# switch gpus for better use when running multiple experiments
if not args.parallel:
torch.cuda.set_device(int(args.gpu))
else:
params.device = torch.device('cpu')
# adapt fase
setting.fase = args.fase
setting.metrics = pd.Series(setting.metrics).drop_duplicates().tolist()
print("metrics {}:".format(setting.metrics))
# Set the random seed for reproducible experiments
torch.manual_seed(230)
if params.cuda: torch.cuda.manual_seed(230)
# Set the logger
logdir=os.path.join(setting_home, setting.fase+"-fase", "runs")
if not args.experiment == '':
logdir=os.path.join(logdir, args.experiment)
if not os.path.isdir(logdir):
os.makedirs(logdir)
# copy params as backupt to logdir
shutil.copy(json_path, os.path.join(logdir, "params.json"))
# utils.set_logger(os.path.join(args.model_dir, 'train.log'))
utils.set_logger(os.path.join(logdir, 'train.log'))
# Create the input data pipeline
logging.info("Loading the datasets...")
# fetch dataloaders
dataloaders = data_loader.fetch_dataloader(args, params, setting, ["train", "valid"])
train_dl = dataloaders['train']
valid_dl = dataloaders['valid']
if setting.num_classes > 1 and params.balance_classes:
train_labels = train_dl.dataset.df[setting.outcome[0]].values
class_weights = compute_class_weight('balanced', np.unique(train_labels), train_labels)
# valid_dl = train_dl
logging.info("- done.")
if args.intercept:
assert len(setting.outcome) == 1, "Multiple outcomes not implemented for intercept yet"
print("running intercept mode")
mu = valid_dl.dataset.df[setting.outcome].values.mean()
def new_forward(self, x, data, mu=mu):
intercept = torch.autograd.Variable(mu * torch.ones((x.shape[0],1)), requires_grad=False).to(params.device, non_blocking=True)
bn_activations = torch.autograd.Variable(torch.zeros((x.shape[0],)), requires_grad=False).to(params.device, non_blocking=True)
return {setting.outcome[0]: intercept, "bn": bn_activations}
net.Net3D.forward = new_forward
params.num_epochs = 1
setting.metrics = []
logdir = os.path.join(logdir, "intercept")
if setting.mode3d:
model = net.Net3D(params, setting).to(params.device)
else:
model = net.CausalNet(params, setting).to(params.device)
optimizers = {'sgd': optim.SGD, 'adam': optim.Adam}
if parallel:
print("parallel mode")
model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
if params.momentum > 0:
optimizer = optimizers[params.optimizer](model.parameters(), lr=params.learning_rate, weight_decay=params.wd, momentum=params.momentum)
else:
optimizer = optimizers[params.optimizer](model.parameters(), lr=params.learning_rate, weight_decay=params.wd)
# if params.use_mi:
# optimizer.add_param_group({'params': mine.parameters()})
if setting.covar_mode and params.lr_t_factor != 1:
optimizer = net.speedup_t(model, params)
if args.restore_last and (not args.cold_start):
print("Loading state dict from last running setting")
utils.load_checkpoint(os.path.join(setting.home, args.fase + "-fase", "last.pth.tar"), model, strict=False)
elif args.restore_warm:
utils.load_checkpoint(os.path.join(setting.home, 'warm-start.pth.tar'), model, strict=False)
else:
pass
# fetch loss function and metrics
if setting.num_classes > 1 and params.balance_classes:
loss_fn = net.get_loss_fn(setting, weights=class_weights)
else:
loss_fn = net.get_loss_fn(setting)
# metrics = {metric:net.all_metrics[metric] for metric in setting.metrics}
metrics = None
if params.monitor_train_tensors:
print(f"Recording all train tensors")
import csv
train_tensor_keys = ['t','x', 'z', 'y', 'x_hat', 'z_hat', 'y_hat']
with open(os.path.join(logdir, 'train-tensors.csv'), 'w') as f:
writer = csv.writer(f)
writer.writerow(['epoch']+train_tensor_keys)
# Train the model
# print(model)
# print(summary(model, (3, 224, 224), batch_size=1))
logging.info("Starting training for {} epoch(s)".format(params.num_epochs))
for split, dl in dataloaders.items():
logging.info("Number of %s samples: %s" % (split, str(len(dl.dataset))))
# logging.info("Number of valid examples: {}".format(len(valid.dataset)))
with SummaryWriter(logdir) as writer:
# train(model, optimizer, loss_fn, train_dl, metrics, params)
train_and_evaluate(model, train_dl, valid_dl, optimizer, loss_fn, metrics, params, setting, args,
writer, logdir, args.restore_file)
Datasets
Models
