"""
Contains functions for running hyperparameter sweep and
Continual Learning model-training and evaluation.
"""
import json
import warnings
from pathlib import Path
from functools import partial
# import random
# import numpy as np
import torch
from ray import tune
from torch import nn, optim
from avalanche.logging import InteractiveLogger, TensorboardLogger
from avalanche.training.plugins import EvaluationPlugin
from avalanche.training.plugins.early_stopping import EarlyStoppingPlugin
from avalanche.evaluation.metrics import (
accuracy_metrics,
loss_metrics,
StreamConfusionMatrix,
)
# Local imports
from utils import models, plotting, data_processing, cl_strategies
from utils.metrics import (
balancedaccuracy_metrics,
sensitivity_metrics,
specificity_metrics,
precision_metrics,
rocauc_metrics,
auprc_metrics,
)
# Suppressing erroneous MaxPool1d named tensors warning
warnings.filterwarnings("once", category=UserWarning)
# GLOBALS
RESULTS_DIR = Path(__file__).parents[1] / "results"
CONFIG_DIR = Path(__file__).parents[1] / "config"
CUDA = torch.cuda.is_available()
DEVICE = "cuda" if CUDA else "cpu"
# Reproducibility
SEED = 12345
# random.seed(SEED)
# np.random.seed(SEED)
torch.manual_seed(SEED)
def save_params(data, domain, outcome, model, strategy, best_params):
"""Save hyper-param config to json."""
file_loc = CONFIG_DIR / data / outcome / domain
file_loc.mkdir(parents=True, exist_ok=True)
with open(
file_loc / f"config_{model}_{strategy}.json", "w", encoding="utf-8"
) as json_file:
json.dump(best_params, json_file)
def load_params(data, domain, outcome, model, strategy):
"""Load hyper-param config from json."""
file_loc = CONFIG_DIR / data / outcome / domain
with open(
file_loc / f"config_{model}_{strategy}.json", encoding="utf-8"
) as json_file:
best_params = json.load(json_file)
return best_params
def save_results(data, outcome, domain, res):
"""Saves results to .json (excluding tensor confusion matrix)."""
with open(
RESULTS_DIR / f"results_{data}_{outcome}_{domain}.json", "w", encoding="utf-8"
) as handle:
res_no_tensors = {
m: {
s: [
{
metric: value
for metric, value in run.items()
if "Confusion" not in metric
}
for run in runs
]
for s, runs in strats.items()
}
for m, strats in res.items()
}
json.dump(res_no_tensors, handle)
def load_strategy(
model,
model_name,
strategy_name,
data="",
domain="",
n_tasks=0,
weight=None,
validate=False,
config=None,
benchmark=None,
early_stopping=False,
):
"""
- `stream` Avg accuracy over all experiences (may rely on tasks being roughly same size?)
- `experience` Accuracy for each experience
"""
strategy = cl_strategies.STRATEGIES[strategy_name]
criterion = nn.CrossEntropyLoss(weight=weight)
if config["generic"]["optimizer"] == "SGD":
optimizer = optim.SGD(
model.parameters(), lr=config["generic"]["lr"], momentum=0.9
)
elif config["generic"]["optimizer"] == "Adam":
optimizer = optim.Adam(model.parameters(), lr=config["generic"]["lr"])
if validate:
loggers = []
else:
timestamp = plotting.get_timestamp()
log_dir = (
RESULTS_DIR
/ "log"
/ "tensorboard"
/ f"{data}_{domain}_{timestamp}"
/ model_name
/ strategy_name
)
interactive_logger = InteractiveLogger()
tb_logger = TensorboardLogger(tb_log_dir=log_dir)
loggers = [interactive_logger, tb_logger]
eval_plugin = EvaluationPlugin(
StreamConfusionMatrix(save_image=False),
loss_metrics(stream=True, experience=not validate),
accuracy_metrics(trained_experience=True, stream=True, experience=not validate),
balancedaccuracy_metrics(
trained_experience=True, stream=True, experience=not validate
),
specificity_metrics(
trained_experience=True, stream=True, experience=not validate
),
sensitivity_metrics(
trained_experience=True, stream=True, experience=not validate
),
precision_metrics(
trained_experience=True, stream=True, experience=not validate
),
# rocauc_metrics(trained_experience=True, stream=True, experience=not validate),
# auprc_metrics(trained_experience=True, stream=True, experience=not validate),
loggers=loggers,
benchmark=benchmark,
)
if early_stopping:
early_stopping = EarlyStoppingPlugin(
patience=5,
val_stream_name="train_stream/Task000",
metric_name="BalancedAccuracy_On_Trained_Experiences",
)
plugins = [early_stopping]
else:
plugins = None
if strategy_name == "Joint":
eval_every = None
cl_strategy = strategy(
model,
optimizer=optimizer,
device=DEVICE,
criterion=criterion,
eval_mb_size=1024,
eval_every=0, # if validate or n_tasks > 5 else 1,
evaluator=eval_plugin,
train_epochs=15,
train_mb_size=config["generic"]["train_mb_size"],
plugins=plugins,
**config["strategy"],
)
return cl_strategy
def train_cl_method(cl_strategy, scenario, strategy_name, validate=False):
"""
Avalanche Cl training loop. For each 'experience' in scenario's train_stream:
- Trains method on experience
- evaluates model on train_stream and test_stream
"""
if not validate:
print("Starting experiment...")
if strategy_name == "Joint":
if not validate:
print(f"Joint training:")
cl_strategy.train(
scenario.train_stream,
eval_streams=[scenario.train_stream, scenario.test_stream],
)
if not validate:
print("Training completed", "\n\n")
else:
for experience in scenario.train_stream:
if not validate:
print(
f"{strategy_name} - Start of experience: {experience.current_experience}"
)
cl_strategy.train(
experience, eval_streams=[scenario.train_stream, scenario.test_stream]
)
if not validate:
print("Training completed", "\n\n")
if validate:
return cl_strategy.evaluator.get_last_metrics()
else:
return cl_strategy.evaluator.get_all_metrics()
def training_loop(
config,
data,
domain,
outcome,
model_name,
strategy_name,
validate=False,
checkpoint_dir=None,
):
"""
Training wrapper:
- loads data
- instantiates model
- equips model with CL strategy
- trains and evaluates method
- returns either results or hyperparam optimisation if `validate`
"""
# Loading data into 'stream' of 'experiences' (tasks)
if not validate:
print("Loading data...")
scenario, n_tasks, n_timesteps, n_channels, weight = data_processing.load_data(
data, domain, outcome, validate
)
if weight is not None:
weight = weight.to(DEVICE)
if not validate:
print("Data loaded.\n")
if not validate:
print(f"N timesteps: {n_timesteps}\nN features: {n_channels}")
model = models.MODELS[model_name](n_channels, n_timesteps, **config["model"])
cl_strategy = load_strategy(
model,
model_name,
strategy_name,
data,
domain,
n_tasks=n_tasks,
weight=weight,
validate=validate,
config=config,
benchmark=scenario,
)
results = train_cl_method(cl_strategy, scenario, strategy_name, validate=validate)
if validate:
loss = results["Loss_Stream/eval_phase/test_stream/Task000"]
accuracy = results[
"Accuracy_On_Trained_Experiences/eval_phase/test_stream/Task000"
]
balancedaccuracy = results[
"BalancedAccuracy_On_Trained_Experiences/eval_phase/test_stream/Task000"
]
# sensitivity = results['Sens_Stream/eval_phase/test_stream/Task000']
# specificity = results['Spec_Stream/eval_phase/test_stream/Task000']
# precision = results['Prec_Stream/eval_phase/test_stream/Task000']
# rocauc = results['ROCAUC_Stream/eval_phase/test_stream/Task000']
# auprc = results['AUPRC_Stream/eval_phase/test_stream/Task000']
# WARNING: `return` overwrites raytune report
tune.report(
loss=loss,
accuracy=accuracy,
balancedaccuracy=balancedaccuracy,
# auprc=auprc,
# rocauc=rocauc
)
else:
return results
def hyperparam_opt(
config, data, domain, outcome, model_name, strategy_name, num_samples
):
"""
Hyperparameter optimisation for the given model/strategy.
Runs over the validation data for the first 2 tasks.
"""
reporter = tune.CLIReporter(
metric_columns=[
"loss",
"accuracy",
"balancedaccuracy",
#'auprc',
#'rocauc'
]
)
resources = {"cpu": 4, "gpu": 0.5} if CUDA else {"cpu": 1}
result = tune.run(
partial(
training_loop,
data=data,
domain=domain,
outcome=outcome,
model_name=model_name,
strategy_name=strategy_name,
validate=True,
),
config=config,
num_samples=num_samples,
progress_reporter=reporter,
raise_on_failed_trial=False,
resources_per_trial=resources,
name=f"{model_name}_{strategy_name}",
local_dir=RESULTS_DIR / "log" / "raytune" / f"{data}_{outcome}_{domain}",
trial_name_creator=lambda t: f"{model_name}_{strategy_name}_{t.trial_id}",
)
best_trial = result.get_best_trial("balancedaccuracy", "max", "last")
print(f"Best trial config: {best_trial.config}")
print(
f"Best trial final validation loss: {best_trial.last_result['loss']}"
)
print(
f"Best trial final validation accuracy: {best_trial.last_result['accuracy']}"
)
print(
f"Best trial final validation balanced accuracy: {best_trial.last_result['balancedaccuracy']}"
)
return best_trial.config
def main(
data,
domain,
outcome,
models,
strategies,
dropout=False,
config_generic={},
config_model={},
config_cl={},
validate=False,
num_samples=50,
freeze_model_hp=False,
):
"""
Main training loop. Defines dataset given outcome/domain
and evaluates model/strategies over given hyperparams over this problem.
"""
# Container for metrics results
res = {m: {s: [] for s in strategies} for m in models}
for model in models:
for strategy in strategies:
# Garbage collection
torch.cuda.empty_cache()
if validate: # Hyperparam opt over first 2 tasks
# Load generic tuned hyper-params
if strategy == "Naive" or not freeze_model_hp:
config = {
"generic": config_generic,
"model": config_model[model],
"strategy": config_cl.get(strategy, {}),
}
else:
naive_params = load_params(data, domain, outcome, model, "Naive")
config = {
"generic": naive_params["generic"],
"model": naive_params["model"],
"strategy": config_cl.get(strategy, {}),
}
# JA: Investigate adding dropout to CNN (final FC layers only?)
if not dropout and model != "CNN":
config["model"]["dropout"] = 0
best_params = hyperparam_opt(
config,
data,
domain,
outcome,
model,
strategy,
num_samples=1 if strategy == "Naive" else num_samples,
)
save_params(data, domain, outcome, model, strategy, best_params)
else: # Training loop over all tasks
config = load_params(data, domain, outcome, model, strategy)
# Multiple runs for Confidence Intervals
n_repeats = 1
for _ in range(n_repeats):
curr_results = training_loop(
config, data, domain, outcome, model, strategy
)
res[model][strategy].append(curr_results)
if not validate:
save_results(data, outcome, domain, res)
plotting.plot_all_figs(data, domain, outcome)
