import numpy as np
import gin
import pytorch_lightning as pl
from models import get_model
import random
import torch
import torch.nn.functional as F
from torch import optim
import math
from utils import Checkpoint, default_device, to_tensor
@gin.configurable
class DeepTimeTrainer(pl.LightningModule):
def __init__(self,
lr,
lambda_lr,
weight_decay,
warmup_epochs,
random_seed,
T_max,
eta_min,
dim_size,
datetime_feats,
):
gin.parse_config_file('/home/reza/Projects/PL_DeepTime/DeepTime/config/config.gin')
super(DeepTimeTrainer, self).__init__()
self.lr = lr
self.lambda_lr = lambda_lr
self.weight_decay = weight_decay
self.warmup_epochs = warmup_epochs
self.random_seed = random_seed
self.lr = lr
self.lambda_lr = lambda_lr
self.weight_decay = weight_decay
self.T_max = T_max
self.warmup_epochs = warmup_epochs
self.eta_min = eta_min
self.model = get_model(
model_type='deeptime',
dim_size=dim_size,
datetime_feats=datetime_feats
)
def on_fit_start(self):
torch.manual_seed(self.random_seed)
np.random.seed(self.random_seed)
random.seed(self.random_seed)
def training_step(self, batch, batch_idx):
x, y, x_time, y_time = map(to_tensor, batch)
forecast = self.model(x, x_time, y_time)
if isinstance(forecast, tuple):
# for models which require reconstruction + forecast loss
loss = F.mse_loss(forecast[0], x) + \
F.mse_loss(forecast[1], y)
else:
loss = F.mse_loss(forecast, y)
self.log('train_loss', loss, prog_bar=True, on_epoch=True)
return {'loss': loss, 'train_loss': loss, }
def training_epoch_end(self, outputs):
avg_train_loss = torch.stack([x["train_loss"] for x in outputs]).mean()
self.log('avg_train_loss', avg_train_loss, on_epoch=True, sync_dist=True)
def validation_step(self, batch, batch_idx):
x, y, x_time, y_time = map(to_tensor, batch)
forecast = self.model(x, x_time, y_time)
if isinstance(forecast, tuple):
# for models which require reconstruction + forecast loss
loss = F.mse_loss(forecast[0], x) + \
F.mse_loss(forecast[1], y)
else:
loss = F.mse_loss(forecast, y)
self.log('val_loss', loss, prog_bar=True, on_epoch=True)
return {'val_loss': loss}
def validation_epoch_end(self, outputs):
return outputs
def test_step(self, batch, batch_idx):
x, y, x_time, y_time = map(to_tensor, batch)
forecast = self.model(x, x_time, y_time)
if isinstance(forecast, tuple):
# for models which require reconstruction + forecast loss
loss = F.mse_loss(forecast[0], x) + \
F.mse_loss(forecast[1], y)
else:
loss = F.mse_loss(forecast, y)
self.log('test_loss', loss, prog_bar=True, on_epoch=True)
return {'test_loss': loss}
def test_epoch_end(self, outputs):
return outputs
@gin.configurable
def configure_optimizers(self):
group1 = [] # lambda
group2 = [] # no decay
group3 = [] # decay
no_decay_list = ('bias', 'norm',)
for param_name, param in self.model.named_parameters():
if '_lambda' in param_name:
group1.append(param)
elif any([mod in param_name for mod in no_decay_list]):
group2.append(param)
else:
group3.append(param)
optimizer = optim.Adam([
{'params': group1, 'weight_decay': 0, 'lr': self.lambda_lr, 'scheduler': 'cosine_annealing'},
{'params': group2, 'weight_decay': 0, 'scheduler': 'cosine_annealing_with_linear_warmup'},
{'params': group3, 'scheduler': 'cosine_annealing_with_linear_warmup'}
], lr=self.lr, weight_decay=self.weight_decay)
scheduler_fns = []
for param_group in optimizer.param_groups:
scheduler = param_group['scheduler']
if scheduler == 'none':
fn = lambda T_cur: 1
elif scheduler == 'cosine_annealing':
lr = eta_max = param_group['lr']
fn = lambda T_cur: (self.eta_min + 0.5 * (eta_max - self.eta_min) * (
1.0 + math.cos(
(T_cur - self.warmup_epochs) / (self.T_max - self.warmup_epochs) * math.pi))) / lr
elif scheduler == 'cosine_annealing_with_linear_warmup':
lr = eta_max = param_group['lr']
fn = lambda T_cur: T_cur / self.warmup_epochs if T_cur < self.warmup_epochs else (self.eta_min + 0.5 * (
eta_max - self.eta_min) * (1.0 + math.cos(
(T_cur - self.warmup_epochs) / (self.T_max - self.warmup_epochs) * math.pi))) / lr
else:
raise ValueError(f'No such scheduler, {scheduler}')
scheduler_fns.append(fn)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=scheduler_fns)
return {'optimizer': optimizer, 'lr_scheduler': scheduler}
def forward(self, batch, z_0=None):
z_0 = None
Y = batch['Y'].to(default_device)
sample_mask = batch['sample_mask'].to(default_device)
available_mask = batch['available_mask'].to(default_device)
# Forecasting
forecasting_mask = available_mask.clone()
if self.n_time_out > 0:
forecasting_mask[:, 0, -self.n_time_out:] = 0
Y, Y_hat, z = self.model(Y=Y, mask=forecasting_mask, idxs=None, z_0=z_0)
if self.n_time_out > 0:
Y = Y[:, :, -self.n_time_out:]
Y_hat = Y_hat[:, :, -self.n_time_out:]
sample_mask = sample_mask[:, :, -self.n_time_out:]
return Y, Y_hat, sample_mask, z
