def test_train_and_resume(self):
train_transformer_style(self.model, self.model_params["training_params"])
self.assertEqual(len(os.listdir("model_save")), 2)
print("first test passed")
model2 = PyTorchForecast(
"MultiAttnHeadSimple",
self.keag_file,
self.keag_file,
self.keag_file,
self.model_params)
data = torch.rand(2, 20, 3)
self.model_params["weight_path"] = os.path.join(
"model_save", sorted(os.listdir("model_save"))[1])
print("Moving to next test")
model3 = PyTorchForecast(
"MultiAttnHeadSimple",
self.keag_file,
self.keag_file,
self.keag_file,
self.model_params)
print("model loaded")
basic_model = model2.model
basic_model.eval()
print("more model stuff")
pre_loaded_model = model3.model
pre_loaded_model.eval()
print("passed model stuff")
self.assertFalse(torch.allclose(pre_loaded_model(data), basic_model(data)))
print("first test good")
self.assertTrue(torch.allclose(basic_model(data), basic_model(data)))
def torch_single_train(model: PyTorchForecast,
opt: optim.Optimizer,
criterion: Type[torch.nn.modules.loss._Loss],
data_loader: DataLoader,
takes_target: bool,
meta_data_model: PyTorchForecast,
meta_data_model_representation: torch.Tensor,
meta_loss=None,
multi_targets=1,
forward_params: Dict = {}) -> float:
print('running torch_single_train')
i = 0
running_loss = 0.0
for src, trg in data_loader:
opt.zero_grad()
# Convert to CPU/GPU/TPU
src = src.to(model.device)
trg = trg.to(model.device)
# TODO figure how to avoid
if meta_data_model:
representation = meta_data_model.model.generate_representation(
meta_data_model_representation)
forward_params["meta_data"] = representation
if meta_loss:
output = meta_data_model.model(meta_data_model_representation)
met_loss = compute_loss(meta_data_model_representation, output,
torch.rand(2, 3, 2), meta_loss, None)
met_loss.backward()
if takes_target:
forward_params["t"] = trg
output = model.model(src, **forward_params)
if multi_targets == 1:
labels = trg[:, :, 0]
elif multi_targets > 1:
labels = trg[:, :, 0:multi_targets]
loss = compute_loss(labels,
output,
src,
criterion,
None,
None,
None,
m=multi_targets)
if loss > 100:
print("Warning: high loss detected")
loss.backward()
opt.step()
if torch.isnan(loss) or loss == float('inf'):
raise ValueError(
"Error infinite or NaN loss detected. Try normalizing data or performing interpolation"
)
running_loss += loss.item()
i += 1
print("The running loss is:")
print(running_loss)
print("The number of items in train is: ")
print(i)
total_loss = running_loss / float(i)
return total_loss
def test_transfer_shit(self):
self.model_params["weight_path"] = os.path.join("model_save", sorted(os.listdir("model_save"))[1])
self.model_params["model_params"]["output_seq_len"] = 6
self.model_params["weight_path_add"] = {}
model3 = PyTorchForecast("MultiAttnHeadSimple", self.keag_file, self.keag_file, self.keag_file,
self.model_params)
# Assert shape is proper
self.assertEqual(2, 2)
data = torch.rand(1, 20, 3)
self.assertEqual(model3.model(data).shape, torch.Size([1, 6]))
def test_model_save(self):
keag_file = os.path.join(self.test_path, "keag_small.csv")
model = PyTorchForecast(
"MultiAttnHeadSimple",
keag_file,
keag_file,
keag_file,
self.model_params)
model.save_model("output", 0)
self.assertEqual(model.training[0][0].shape, torch.Size([20, 3]))
def setUp(self):
self.test_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_init")
self.model_params = {"model_params": {"number_time_series": 3, "seq_len": 20},
"dataset_params": {"forecast_history": 20, "class": "default", "forecast_length": 20,
"relevant_cols": ["cfs", "temp", "precip"], "target_col": ["cfs"],
"interpolate": False},
"training_params": {"optimizer": "Adam", "lr": .1, "criterion": "MSE", "epochs": 1,
"batch_size": 2, "optim_params": {}},
"wandb": False, "inference_params": {"hours_to_forecast": 10}}
self.keag_file = os.path.join(self.test_path, "keag_small.csv")
self.model = PyTorchForecast("MultiAttnHeadSimple", self.keag_file, self.keag_file, self.keag_file,
self.model_params)
self.dummy_model = PyTorchForecast("DummyTorchModel", self.keag_file, self.keag_file, self.keag_file,
{"model_params": {"forecast_length": 5},
"dataset_params": {"forecast_history": 5, "class": "default",
"forecast_length": 5,
"relevant_cols": ["cfs", "temp", "precip"],
"target_col": ["cfs"], "interpolate": False,
"train_end": 100},
"training_params": {"optimizer": "Adam", "lr": .1, "criterion": "MSE",
"epochs": 1, "batch_size": 2, "optim_params": {}},
"inference_params": {"hours_to_forecast": 15},
"wandb": False})
self.full_transformer_params = {"use_decoder": True, "model_params": {"number_time_series": 3, "seq_length": 20,
"output_seq_len": 15},
"dataset_params": {"forecast_history": 20, "class": "default",
"forecast_length": 15,
"relevant_cols": ["cfs", "temp", "precip"],
"target_col": ["cfs"], "interpolate": False, "train_end": 50,
"valid_end": 100},
"training_params": {"optimizer": "Adam", "lr": .01, "criterion": "MSE",
"epochs": 1, "batch_size": 2, "optim_params": {}},
"inference_params": {"hours_to_forecast": 10},
"wandb": False}
self.simple_param = {"use_decoder": True,
"model_params": {"n_time_series": 3, "seq_length": 80, "output_seq_len": 20},
"dataset_params": {"forecast_history": 20, "class": "default", "forecast_length": 15,
"relevant_cols": ["cfs", "temp", "precip"], "target_col": ["cfs"],
"interpolate": False, "train_end": 50,
"valid_end": 100}, "inference_params": {"hours_to_forecast": 10},
"training_params": {"optimizer": "Adam", "lr": .01, "criterion": "MSE", "epochs": 1,
"batch_size": 2, "optim_params": {}},
"wandb": False}
self.transformer = PyTorchForecast("SimpleTransformer", self.keag_file, self.keag_file, self.keag_file,
self.full_transformer_params)
self.simple_linear_model = PyTorchForecast("SimpleLinearModel", self.keag_file, self.keag_file, self.keag_file,
self.simple_param)
self.opt = torch.optim.Adam(self.dummy_model.model.parameters(), lr=0.0001)
self.criterion = torch.nn.modules.loss.MSELoss()
self.data_loader = DataLoader(self.dummy_model.training, batch_size=2, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None)
def test_pytorch_wrapper_default(self):
keag_file = os.path.join(self.test_path, "keag_small.csv")
model = PyTorchForecast("MultiAttnHeadSimple", keag_file, keag_file,
keag_file, self.model_params)
self.assertEqual(model.model.dense_shape.in_features, 3)
self.assertEqual(model.model.multi_attn.embed_dim, 128)
self.assertEqual(model.model.multi_attn.num_heads, 8)
def torch_single_train(model: PyTorchForecast,
opt: optim.Optimizer,
criterion: Type[torch.nn.modules.loss._Loss],
data_loader: DataLoader,
takes_target: bool,
forward_params: Dict = {}) -> float:
i = 0
running_loss = 0.0
for src, trg in data_loader:
opt.zero_grad()
# Convert to CPU/GPU/TPU
src = src.to(model.device)
trg = trg.to(model.device)
# TODO figure how to avoid
if takes_target:
forward_params["t"] = trg
output = model.model(src, **forward_params)
labels = trg[:, :, 0]
loss = criterion(output, labels.float())
if loss > 100:
print("Warning: high loss detected")
loss.backward()
opt.step()
if torch.isnan(loss) or loss == float('inf'):
raise (
"Error infinite or NaN loss detected. Try normalizing data or performing interpolation"
)
running_loss += loss.item()
i += 1
print("The running loss is:")
print(running_loss)
print("The number of items in train is: ")
print(i)
total_loss = running_loss / float(i)
return total_loss
def test_informer_init(self):
import json
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_informer.json")) as y:
json_params = json.load(y)
keag_file = os.path.join(self.test_path, "keag_small.csv")
inf = PyTorchForecast("Informer", keag_file, keag_file, keag_file, json_params)
self.assertTrue(inf)
def load_model(model_params_dict, file_path, weight_path: str) -> PyTorchForecast:
if weight_path:
model_params_dict["weight_path"] = weight_path
model_params_dict["inference_params"]["test_csv_path"] = file_path
model_params_dict["inference_params"]["dataset_params"]["file_path"] = file_path
m = PyTorchForecast(model_params_dict["model_name"], file_path, file_path, file_path, model_params_dict)
return m
def test_handle_meta(self):
with open(os.path.join(os.path.dirname(__file__), "da_meta.json")) as f:
json_config = json.load(f)
model = PyTorchForecast("DARNN", self.keag_file, self.keag_file, self.keag_file, json_config)
meta_models, meta_reps, loss = handle_meta_data(model)
self.assertIsNone(loss)
self.assertIsInstance(meta_reps, torch.Tensor)
self.assertIsInstance(meta_models, PyTorchForecast)
def train_function(model_type: str, params:Dict):
"""
Function to train a Model(TimeSeriesModel) or da_rnn. Will return the trained model
model_type str: Type of the model (for now) must be da_rnn or
:params dict: Dictionary containing all the parameters needed to run the model
"""
dataset_params = params["dataset_params"]
if model_type == "da_rnn":
from flood_forecast.da_rnn.train_da import da_rnn, train
from flood_forecast.preprocessing.preprocess_da_rnn import make_data
preprocessed_data = make_data(params["dataset_params"]["training_path"], params["dataset_params"]["target_col"], params["dataset_params"]["forecast_length"])
config, model = da_rnn(preprocessed_data, len(dataset_params["target_col"]))
# All train functions return trained_model
trained_model = train(model, preprocessed_data, config)
elif model_type == "PyTorch":
trained_model = PyTorchForecast(
params["model_name"],
dataset_params["training_path"],
dataset_params["validation_path"],
dataset_params["test_path"],
params)
train_transformer_style(trained_model, params["training_params"], params["forward_params"])
params["inference_params"]["dataset_params"]["scaling"] = scaler_dict[dataset_params["scaler"]]
test_acc = evaluate_model(
trained_model,
model_type,
params["dataset_params"]["target_col"],
params["metrics"],
params["inference_params"],
{})
wandb.run.summary["test_accuracy"] = test_acc[0]
df_train_and_test = test_acc[1]
forecast_start_idx = test_acc[2]
df_prediction_samples = test_acc[3]
inverse_mae = 1 / (
df_train_and_test.loc[forecast_start_idx:, "preds"] -
df_train_and_test.loc[forecast_start_idx:, params["dataset_params"]["target_col"][0]]).abs()
pred_std = df_prediction_samples.std(axis=1)
average_prediction_sharpe = (inverse_mae / pred_std).mean()
wandb.log({'average_prediction_sharpe': average_prediction_sharpe})
# Log plots
test_plot = plot_df_test_with_confidence_interval(
df_train_and_test,
df_prediction_samples,
forecast_start_idx,
params,
ci=95,
alpha=0.25)
wandb.log({"test_plot": test_plot})
test_plot_all = go.Figure()
for relevant_col in params["dataset_params"]["relevant_cols"]:
test_plot_all.add_trace(go.Scatter(x=df_train_and_test.index, y=df_train_and_test[relevant_col], name=relevant_col))
wandb.log({"test_plot_all": test_plot_all})
else:
raise Exception("Please supply valid model type for forecasting")
return trained_model
def test_data_correct(self):
keag_file = os.path.join(self.test_path, "keag_small.csv")
model = PyTorchForecast(
"MultiAttnHeadSimple",
keag_file,
keag_file,
keag_file,
self.model_params)
model
def test_pytorch_wrapper_custom(self):
self.model_params["model_params"] = {
"number_time_series": 6,
"d_model": 112
}
keag_file = os.path.join(self.test_path, "keag_small.csv")
model = PyTorchForecast("MultiAttnHeadSimple", keag_file, keag_file,
keag_file, self.model_params)
self.assertEqual(model.model.dense_shape.in_features, 6)
self.assertEqual(model.model.multi_attn.embed_dim, 112)
def test_simple_transformer(self):
self.model_params["model_params"] = {
"seq_length": 19,
"number_time_series": 6,
"d_model": 136,
"n_heads": 8
}
keag_file = os.path.join(self.test_path, "keag_small.csv")
model = PyTorchForecast("SimpleTransformer", keag_file, keag_file,
keag_file, self.model_params)
self.assertEqual(model.model.dense_shape.in_features, 6)
self.assertEqual(model.model.mask.shape, torch.Size([19, 19]))
def test_removing_layer_param(self):
model3 = PyTorchForecast("MultiAttnHeadSimple", self.keag_file, self.keag_file, self.keag_file,
self.model_params)
model3.save_model("output.pth", 2)
self.model_params["model_params"]["output_seq_len"] = 7
self.model_params["weight_path_add"] = {}
self.model_params["weight_path_add"]["excluded_layers"] = ["last_layer.weight", "last_layer.bias"]
model = PyTorchForecast("MultiAttnHeadSimple", self.keag_file, self.keag_file, self.keag_file,
self.model_params)
result = model.model(torch.rand(1, 20, 3))
self.assertEqual(result.shape[1], 7)
def convert_to_torch_script(model: PyTorchForecast, save_path: str) -> PyTorchForecast:
"""Function to convert PyTorch model to torch script and save
:param model: The PyTorchForecast model you wish to convert
:type model: PyTorchForecast
:param save_path: File name to save the TorchScript model under.
:type save_path: str
:return: Returns the model with an added .script_model attribute
:rtype: PyTorchForecast
"""
model.model.eval()
forecast_history = model.params["dataset_params"]["forecast_history"]
n_features = model.params["model_params"]["n_time_series"]
test_input = torch.rand(2, forecast_history, n_features)
model_script = torch.jit.trace(model.model, test_input)
test_input1 = torch.rand(4, forecast_history, n_features)
a = model_script(test_input1)
b = model.model(test_input1)
model.script_model = model_script
assert torch.eq(a, b).all()
model_script.save(save_path)
return model
class ModelInterpretabilityTest(unittest.TestCase):
test_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "test_init"
)
test_path2 = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "test_data"
)
model_params = {
"model_params": {
"number_time_series": 3,
"seq_len": 20,
"output_seq_len": 10,
},
"dataset_params": {
"forecast_history": 20,
"class": "default",
"forecast_length": 20,
"relevant_cols": ["cfs", "temp", "precip"],
"target_col": ["cfs"],
"interpolate": False,
},
"wandb": {
"name": "flood_forecast_circleci",
"tags": ["dummy_run", "circleci"],
"project": "repo-flood_forecast",
},
"inference_params": {"hours_to_forecast": 50},
}
keag_file = os.path.join(test_path, "keag_small.csv")
model = PyTorchForecast(
"MultiAttnHeadSimple", keag_file, keag_file, keag_file, model_params,
)
data_base_params = {
"file_path": os.path.join(test_path2, "keag_small.csv"),
"forecast_history": 20,
"forecast_length": 20,
"relevant_cols": ["cfs", "temp", "precip"],
"target_col": ["cfs"],
"interpolate_param": False,
}
def test_deep_explain_model_summary_plot(self):
deep_explain_model_summary_plot(
model=self.model,
datetime_start=datetime.datetime(2014, 6, 2, 0),
test_csv_path=os.path.join(self.test_path2, "keag_small.csv"),
dataset_params=self.data_base_params,
)
# dummy assert
self.assertEqual(1, 1)
def train_function(model_type: str, params: Dict) -> PyTorchForecast:
params["forward_params"] = {}
dataset_params = params["dataset_params"]
dataset_params["forecast_history"] = 1
dataset_params["forecast_length"] = 1
dataset_params["forecast_something_or"] = 1
trained_model = PyTorchForecast(
params["model_name"],
dataset_params["training_path"],
dataset_params["validation_path"],
dataset_params["test_path"],
params)
train_transformer_style(trained_model, params["training_params"], params["forward_params"])
return trained_model
def train_function(model_type: str, params: Dict) -> PyTorchForecast:
""" Function to train meta data-models"""
params["forward_params"] = {}
dataset_params = params["dataset_params"]
if "forecast_history" not in dataset_params:
dataset_params["forecast_history"] = 1
dataset_params["forecast_length"] = 1
trained_model = PyTorchForecast(params["model_name"],
dataset_params["training_path"],
dataset_params["validation_path"],
dataset_params["test_path"], params)
train_transformer_style(trained_model, params["training_params"],
params["forward_params"])
return trained_model
def handle_meta_data(model):
meta_loss = None
with open(model.params["meta_data"]["path"]) as f:
json_data = json.load(f)
if "meta_loss" in model.params["meta_data"]:
meta_loss_str = model.params["meta_data"]["meta_loss"]
meta_loss = pytorch_criterion_dict[meta_loss_str]()
dataset_params2 = json_data["dataset_params"]
training_path = dataset_params2["training_path"]
valid_path = dataset_params2["validation_path"]
meta_name = json_data["model_name"]
meta_model = PyTorchForecast(meta_name, training_path, valid_path,
dataset_params2["test_path"], json_data)
meta_representation = get_meta_representation(
model.params["meta_data"]["column_id"],
model.params["meta_data"]["uuid"], meta_model)
return meta_model, meta_representation, meta_loss
def torch_single_train(model: PyTorchForecast,
opt: optim.Optimizer,
criterion: Type[torch.nn.modules.loss._Loss],
data_loader: DataLoader,
takes_target: bool,
meta_data_model: PyTorchForecast,
meta_data_model_representation: torch.Tensor,
forward_params: Dict = {}) -> float:
i = 0
running_loss = 0.0
for src, trg in data_loader:
opt.zero_grad()
# Convert to CPU/GPU/TPU
src = src.to(model.device)
trg = trg.to(model.device)
# TODO figure how to avoid
if meta_data_model:
representation = meta_data_model.model.generate_representation(
meta_data_model_representation)
forward_params["meta_data"] = representation
if takes_target:
forward_params["t"] = trg
output = model.model(src, **forward_params)
labels = trg[:, :, 0]
if isinstance(criterion, GaussianLoss):
g_loss = GaussianLoss(output[0], output[1])
loss = g_loss(labels)
else:
loss = criterion(output, labels.float())
# TODO fix Guassian loss
if loss > 100:
print("Warning: high loss detected")
loss.backward()
opt.step()
if torch.isnan(loss) or loss == float('inf'):
raise ValueError(
"Error infinite or NaN loss detected. Try normalizing data or performing interpolation"
)
running_loss += loss.item()
i += 1
print("The running loss is:")
print(running_loss)
print("The number of items in train is: ")
print(i)
total_loss = running_loss / float(i)
return total_loss
def load_model(model_params_dict, file_path: str, weight_path: str) -> PyTorchForecast:
"""Function to load a PyTorchForecast model from an existing config file.
:param model_params_dict: Dictionary of model parameters
:type model_params_dict: Dict
:param file_path: [description]
:type file_path: str
:param weight_path: [description]
:type weight_path: str
:return: [description]
:rtype: PyTorchForecast
"""
if weight_path:
model_params_dict["weight_path"] = weight_path
model_params_dict["inference_params"]["test_csv_path"] = file_path
model_params_dict["inference_params"]["dataset_params"]["file_path"] = file_path
m = PyTorchForecast(model_params_dict["model_name"], file_path, file_path, file_path, model_params_dict)
return m
def handle_meta_data(model: PyTorchForecast):
"""A function to initialize models with meta-data
:param model: A PyTorchForecast model with meta_data parameter block in config file.
:type model: PyTorchForecast
:return: Returns a tuple of the initial meta-representation
:rtype: tuple(PyTorchForecast, torch.Tensor, float)
"""
meta_loss = None
with open(model.params["meta_data"]["path"]) as f:
json_data = json.load(f)
if "meta_loss" in model.params["meta_data"]:
meta_loss_str = model.params["meta_data"]["meta_loss"]
meta_loss = pytorch_criterion_dict[meta_loss_str]()
dataset_params2 = json_data["dataset_params"]
training_path = dataset_params2["training_path"]
valid_path = dataset_params2["validation_path"]
meta_name = json_data["model_name"]
meta_model = PyTorchForecast(meta_name, training_path, valid_path, dataset_params2["test_path"], json_data)
meta_representation = get_meta_representation(model.params["meta_data"]["column_id"],
model.params["meta_data"]["uuid"], meta_model)
return meta_model, meta_representation, meta_loss
def load_model(model_params_dict, file_path: str,
weight_path: str) -> PyTorchForecast:
"""Function to load a PyTorchForecast model from an existing config file.
:param model_params_dict: Dictionary of model parameters
:type model_params_dict: Dict
:param file_path: The path to the CSV for running infer
:type file_path: str
:param weight_path: The path to the model weights (can be GCS)
:type weight_path: str
:return: Returns a PyTorchForecast model initialized with the proper data
:rtype: PyTorchForecast
"""
if weight_path:
model_params_dict["weight_path"] = weight_path
model_params_dict["inference_params"]["test_csv_path"] = file_path
model_params_dict["inference_params"]["dataset_params"][
"file_path"] = file_path
if "weight_path_add" in model_params_dict:
if "excluded_layers" in model_params_dict["weight_path_add"]:
del model_params_dict["weight_path_add"]["excluded_layers"]
m = PyTorchForecast(model_params_dict["model_name"], file_path, file_path,
file_path, model_params_dict)
return m
def train_transformer_style(model: PyTorchForecast,
training_params: Dict,
takes_target=False,
forward_params: Dict = {},
model_filepath: str = "model_save") -> None:
"""Function to train any PyTorchForecast model
:param model: A properly wrapped PyTorchForecast model
:type model: PyTorchForecast
:param training_params: A dictionary of the necessary parameters for training.
:type training_params: Dict
:param takes_target: A parameter to determine whether a model requires the target, defaults to False
:type takes_target: bool, optional
:param forward_params: [description], defaults to {}
:type forward_params: Dict, optional
:param model_filepath: The file path to load modeel weights from, defaults to "model_save"
:type model_filepath: str, optional
:raises ValueError: Has an error
"""
use_wandb = model.wandb
es = None
worker_num = 1
pin_memory = False
dataset_params = model.params["dataset_params"]
num_targets = 1
if "n_targets" in model.params:
num_targets = model.params["n_targets"]
if "num_workers" in dataset_params:
worker_num = dataset_params["num_workers"]
print("using " + str(worker_num))
if "pin_memory" in dataset_params:
pin_memory = dataset_params["pin_memory"]
print("Pin memory set to true")
if "early_stopping" in model.params:
es = EarlyStopper(model.params["early_stopping"]['patience'])
opt = pytorch_opt_dict[training_params["optimizer"]](
model.model.parameters(), **training_params["optim_params"])
criterion_init_params = {}
if "criterion_params" in training_params:
criterion_init_params = training_params["criterion_params"]
criterion = pytorch_criterion_dict[training_params["criterion"]](
**criterion_init_params)
if "probabilistic" in model.params[
"model_params"] or "probabilistic" in model.params:
probabilistic = True
else:
probabilistic = False
max_epochs = training_params["epochs"]
data_loader = DataLoader(model.training,
batch_size=training_params["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=worker_num,
collate_fn=None,
pin_memory=pin_memory,
drop_last=False,
timeout=0,
worker_init_fn=None)
validation_data_loader = DataLoader(
model.validation,
batch_size=training_params["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=worker_num,
collate_fn=None,
pin_memory=pin_memory,
drop_last=False,
timeout=0,
worker_init_fn=None)
# TODO support batch_size > 1
test_data_loader = DataLoader(model.test_data,
batch_size=1,
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=worker_num,
collate_fn=None,
pin_memory=pin_memory,
drop_last=False,
timeout=0,
worker_init_fn=None)
meta_model = None
meta_representation = None
meta_loss = None
if model.params.get("meta_data") is None:
model.params["meta_data"] = False
if model.params["meta_data"]:
meta_model, meta_representation, meta_loss = handle_meta_data(model)
if use_wandb:
wandb.watch(model.model)
use_decoder = False
if "use_decoder" in model.params:
use_decoder = True
session_params = []
for epoch in range(max_epochs):
total_loss = torch_single_train(model,
opt,
criterion,
data_loader,
takes_target,
meta_model,
meta_representation,
meta_loss,
multi_targets=num_targets,
forward_params=forward_params.copy())
print("The loss for epoch " + str(epoch))
print(total_loss)
valid = compute_validation(
validation_data_loader,
model.model,
epoch,
model.params["dataset_params"]["forecast_length"],
model.crit,
model.device,
multi_targets=num_targets,
meta_model=meta_model,
decoder_structure=use_decoder,
use_wandb=use_wandb,
probabilistic=probabilistic)
if valid == 0.0:
raise ValueError(
"Error validation loss is zero there is a problem with the validator."
)
if use_wandb:
wandb.log({'epoch': epoch, 'loss': total_loss})
epoch_params = {
"epoch": epoch,
"train_loss": str(total_loss),
"validation_loss": str(valid)
}
session_params.append(epoch_params)
if es:
if not es.check_loss(model.model, valid):
print("Stopping model now")
model.model.load_state_dict(torch.load("checkpoint.pth"))
break
decoder_structure = True
if model.params["dataset_params"]["class"] == "AutoEncoder":
decoder_structure = False
test = compute_validation(
test_data_loader,
model.model,
epoch,
model.params["dataset_params"]["forecast_length"],
model.crit,
model.device,
meta_model=meta_model,
multi_targets=num_targets,
decoder_structure=decoder_structure,
use_wandb=use_wandb,
val_or_test="test_loss",
probabilistic=probabilistic)
print("test loss:", test)
model.params["run"] = session_params
model.save_model(model_filepath, max_epochs)
def torch_single_train(model: PyTorchForecast,
opt: optim.Optimizer,
criterion: Type[torch.nn.modules.loss._Loss],
data_loader: DataLoader,
takes_target: bool,
meta_data_model: PyTorchForecast,
meta_data_model_representation: torch.Tensor,
meta_loss=None,
multi_targets=1,
forward_params: Dict = {}) -> float:
probablistic = None
if "probabilistic" in model.params["model_params"]:
probablistic = True
print('running torch_single_train')
i = 0
output_std = None
mulit_targets_copy = multi_targets
running_loss = 0.0
for src, trg in data_loader:
opt.zero_grad()
# Convert to CPU/GPU/TPU
if meta_data_model:
representation = meta_data_model.model.generate_representation(
meta_data_model_representation)
forward_params["meta_data"] = representation
if meta_loss:
output = meta_data_model.model(meta_data_model_representation)
met_loss = compute_loss(meta_data_model_representation, output,
torch.rand(2, 3, 2), meta_loss, None)
met_loss.backward()
if takes_target:
forward_params["t"] = trg
elif "TemporalLoader" == model.params["dataset_params"]["class"]:
forward_params["x_mark_enc"] = src[1].to(model.device)
forward_params["x_dec"] = trg[1].to(model.device)
forward_params["x_mark_dec"] = trg[0].to(model.device)
src = src[0]
# Assign to avoid other if statement
trg = trg[0]
src = src.to(model.device)
trg = trg.to(model.device)
output = model.model(src, **forward_params)
if hasattr(model.model, "pred_len"):
multi_targets = mulit_targets_copy
pred_len = model.model.pred_len
labels = trg[:, -pred_len:, 0:multi_targets]
multi_targets = False
if multi_targets == 1:
labels = trg[:, :, 0]
elif multi_targets > 1:
labels = trg[:, :, 0:multi_targets]
if probablistic:
output1 = output
output = output.mean
output_std = output1.stddev
loss = compute_loss(labels,
output,
src,
criterion,
None,
probablistic,
output_std,
m=multi_targets)
if loss > 100:
print("Warning: high loss detected")
loss.backward()
opt.step()
if torch.isnan(loss) or loss == float('inf'):
raise ValueError(
"Error infinite or NaN loss detected. Try normalizing data or performing interpolation"
)
running_loss += loss.item()
i += 1
print("The running loss is: ")
print(running_loss)
print("The number of items in train is: " + str(i))
total_loss = running_loss / float(i)
return total_loss
def torch_single_train(model: PyTorchForecast,
opt: optim.Optimizer,
criterion: Type[torch.nn.modules.loss._Loss],
data_loader: DataLoader,
takes_target: bool,
meta_data_model: PyTorchForecast,
meta_data_model_representation: torch.Tensor,
meta_loss=None,
multi_targets=1,
forward_params: Dict = {}) -> float:
"""Function that performs training of a single model. Runs through one epoch of the data.
:param model: The PyTorchForecast model that is trained
:type model: PyTorchForecast
:param opt: The optimizer to use in the code
:type opt: optim.Optimizer
:param criterion: [description]
:type criterion: Type[torch.nn.modules.loss._Loss]
:param data_loader: [description]
:type data_loader: DataLoader
:param takes_target: A boolean that indicates whether the model takes the target during training
:type takes_target: bool
:param meta_data_model: If supplied a model that handles meta-data else None.
:type meta_data_model: PyTorchForecast
:param meta_data_model_representation: [description]
:type meta_data_model_representation: torch.Tensor
:param meta_loss: [description], defaults to None
:type meta_loss: [type], optional
:param multi_targets: [description], defaults to 1
:type multi_targets: int, optional
:param forward_params: [description], defaults to {}
:type forward_params: Dict, optional
:raises ValueError: [description]
:return: [description]
:rtype: float
"""
probablistic = None
if "probabilistic" in model.params["model_params"]:
probablistic = True
print('running torch_single_train')
i = 0
output_std = None
mulit_targets_copy = multi_targets
running_loss = 0.0
for src, trg in data_loader:
opt.zero_grad()
if meta_data_model:
representation = meta_data_model.model.generate_representation(meta_data_model_representation)
forward_params["meta_data"] = representation
if meta_loss:
output = meta_data_model.model(meta_data_model_representation)
met_loss = compute_loss(meta_data_model_representation, output, torch.rand(2, 3, 2), meta_loss, None)
met_loss.backward()
if takes_target:
forward_params["t"] = trg
elif "TemporalLoader" == model.params["dataset_params"]["class"]:
forward_params["x_mark_enc"] = src[1].to(model.device)
forward_params["x_dec"] = trg[1].to(model.device)
forward_params["x_mark_dec"] = trg[0].to(model.device)
src = src[0]
pred_len = model.model.pred_len
trg = trg[0]
trg[:, -pred_len:, :] = torch.zeros_like(trg[:, -pred_len:, :].long()).float().to(model.device)
# Assign to avoid other if statement
elif "SeriesIDLoader" == model.params["dataset_params"]["class"]:
pass
src = src.to(model.device)
trg = trg.to(model.device)
output = model.model(src, **forward_params)
if hasattr(model.model, "pred_len"):
multi_targets = mulit_targets_copy
pred_len = model.model.pred_len
output = output[:, :, 0:multi_targets]
labels = trg[:, -pred_len:, 0:multi_targets]
multi_targets = False
if model.params["dataset_params"]["class"] == "GeneralClassificationLoader":
labels = trg
elif multi_targets == 1:
labels = trg[:, :, 0]
elif multi_targets > 1:
labels = trg[:, :, 0:multi_targets]
if probablistic:
output1 = output
output = output.mean
output_std = output1.stddev
if type(criterion) == list:
loss = multi_crit(criterion, output, labels, None)
else:
loss = compute_loss(labels, output, src, criterion, None, probablistic, output_std, m=multi_targets)
if loss > 100:
print("Warning: high loss detected")
loss.backward()
opt.step()
if torch.isnan(loss) or loss == float('inf'):
raise ValueError("Error infinite or NaN loss detected. Try normalizing data or performing interpolation")
running_loss += loss.item()
i += 1
print("The running loss is: ")
print(running_loss)
print("The number of items in train is: " + str(i))
total_loss = running_loss / float(i)
return total_loss
def test_ae(self):
model = PyTorchForecast("DARNN", self.keag_file, self.keag_file, self.keag_file, self.meta_model_params)
for parameter in model.model.parameters():
self.assertTrue(parameter.requires_grad)
def train_transformer_style(model: PyTorchForecast,
training_params: Dict,
takes_target=False,
forward_params: Dict = {},
model_filepath: str = "model_save") -> None:
"""
Function to train any PyTorchForecast model
:model The initialized PyTorchForecastModel
:training_params_dict A dictionary of the parameters needed to train model
:takes_target boolean: Determines whether to pass target during training
:forward_params: A dictionary for additional forward parameters (for instance target)
"""
use_wandb = model.wandb
es = None
if "early_stopping" in model.params:
es = EarlyStopper(model.params["early_stopping"]['patience'])
opt = pytorch_opt_dict[training_params["optimizer"]](
model.model.parameters(), **training_params["optim_params"])
criterion_init_params = {}
if "criterion_params" in training_params:
criterion_init_params = training_params["criterion_params"]
criterion = pytorch_criterion_dict[training_params["criterion"]](
**criterion_init_params)
max_epochs = training_params["epochs"]
data_loader = DataLoader(model.training,
batch_size=training_params["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
timeout=0,
worker_init_fn=None)
validation_data_loader = DataLoader(
model.validation,
batch_size=training_params["batch_size"],
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
timeout=0,
worker_init_fn=None)
test_data_loader = DataLoader(model.test_data,
batch_size=1,
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
timeout=0,
worker_init_fn=None)
if use_wandb:
import wandb
wandb.watch(model.model)
session_params = []
for epoch in range(max_epochs):
total_loss = torch_single_train(model, opt, criterion, data_loader,
takes_target, forward_params)
print("The loss for epoch " + str(epoch))
print(total_loss)
use_decoder = False
if "use_decoder" in model.params:
use_decoder = True
valid = compute_validation(
validation_data_loader,
model.model,
epoch,
model.params["dataset_params"]["forecast_length"],
criterion,
model.device,
decoder_structure=use_decoder,
use_wandb=use_wandb)
if valid < 0.01:
raise (
"Error validation loss is zero there is a problem with the validator."
)
if use_wandb:
wandb.log({'epoch': epoch, 'loss': total_loss})
epoch_params = {
"epoch": epoch,
"train_loss": str(total_loss),
"validation_loss": str(valid)
}
session_params.append(epoch_params)
if es:
if not es.check_loss(model.model, valid):
print("Stopping model now")
model.model.load_state_dict(torch.load("checkpoint.pth"))
break
decoder_structure = True
if model.params["dataset_params"]["class"] != "default":
decoder_structure = False
test = compute_validation(
test_data_loader,
model.model,
epoch,
model.params["dataset_params"]["forecast_length"],
criterion,
model.device,
decoder_structure=decoder_structure,
use_wandb=use_wandb,
val_or_test="test_loss")
print("test loss:", test)
model.params["run"] = session_params
model.save_model(model_filepath, max_epochs)
def load_model(model_params_dict, file_path, weight_path) -> PyTorchForecast:
if weight_path:
model_params_dict["weight_path"] = weight_path
m = PyTorchForecast(model_params_dict["model_name"], file_path, file_path,
file_path, model_params_dict)
return m
