aqacf-ml-main/firealarm_ml/epa_seq2seq.py

#!/usr/bin/env python
"""
EPA Sequence-to-Sequence Model
see https://sladewinter.medium.com/video-frame-prediction-using-convlstm-network-in-pytorch-b5210a6ce582
"""
import os
import json
import click
import numpy as np
import xarray as xr
from tqdm import tqdm
import torch
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import dagshub
import mlflow

from epa_seq2seq_model import EpaSeq2Seq
dagshub.init("aqacf-ml-main", "ML-Purdue")


class EpaDataset(Dataset):


    def __init__(self, dataset, window_size, in_channels, out_channels):
        self.dataset = dataset
        self.X = self.dataset[in_channels]
        self.y = self.dataset[out_channels]
        self.times = dataset.time.values
        self.window_size = window_size
        self._n = len(self.times) - (self.window_size - 1)


    @property
    def frame_size(self):
        return (
            len(self.dataset.lat.values),
            len(self.dataset.lon.values),
        )


    def __len__(self):
        return self._n


    def get_time(self, i):
        return self.times[i: i + self.window_size]


    def __getitem__(self, i):
        X = self.X.isel(time=slice(i, i + self.window_size))
        y = self.y.isel(time=slice(i, i + self.window_size))

        X = np.transpose(
            np.stack([X[v] for v in X.data_vars]),
            (0, 1, 2, 3)
        )

        y = np.transpose(
            np.stack([y[v] for v in y.data_vars]),
            (0, 1, 2, 3)
        )

        return X, y


def mse_nan_loss(output, target):
    mask = ~torch.isnan(target)
    sqe = (target[mask] - output[mask])**2
    return sqe.mean()


def train_model(data_loader, model, loss_function, optimizer):
    num_batches = len(data_loader)
    total_loss = 0
    model.train()
    
    n = 0
    for X, y in tqdm(data_loader, 'Training', total=num_batches):
        X, y = X.cuda(), y.cuda()
        output = model(X)
        loss = loss_function(output, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

        n += 1
        #if n > 2: break

    #avg_loss = total_loss / num_batches
    avg_loss = total_loss / n
    rmse = np.sqrt(avg_loss)
    print(f"Train loss: {rmse}")
    return avg_loss


def eval_model(data_loader, model, loss_function):
    
    num_batches = len(data_loader)
    total_loss = 0

    model.eval()
    with torch.no_grad():
        n = 0
        for X, y in tqdm(data_loader, 'Evaluating', total=num_batches):
            X, y = X.cuda(), y.cuda()
            output = model(X)
            total_loss += loss_function(output, y).item()
            n += 1
            #if n > 2: break

    #avg_loss = total_loss / num_batches
    avg_loss = total_loss / n
    rmse = np.sqrt(avg_loss)
    print(f"Validation loss: {rmse}")
    return avg_loss


def save_state(outputdir, config, epoch, model, optimizer,
        train_loss, validation_loss, loss_function_name):

    state_dict = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'train_loss': train_loss,
        'validation_loss': validation_loss,
        'loss_function_name': loss_function_name,
        'config': config,
    }
    path = os.path.join(outputdir, f'state_epoch_{epoch:05d}.pt')
    torch.save(state_dict, path)


@click.command()
@click.argument('configfile')
@click.argument('dropoutconfigfile')
@click.argument('datafile')
@click.argument('outputdir')
def main(configfile, dropoutconfigfile, datafile, outputdir):

    with open(configfile, 'r') as f:
        config = json.load(f)

    checkpoint_path = r"C:\Users\Michael\aqacf-ml-main\data\train_output\train_transfer_step1_trial1\state_epoch_00020.pt"

    window_size = config['sequence_length']
    window_size = 1
    in_channels = config['in_channels']
    out_channels = config['out_channels']
    batch_size = config['batch_size']
    epochs = config['epochs']
    learning_rate = config['learning_rate']
    model_params = config['model_params']

    with open(dropoutconfigfile, 'r') as f:
        dropoutConfig = json.load(f)

    dataset_names = ('train', 'validation', 'test')
    datasets = {}

    with xr.open_zarr(datafile) as ds:
        for s in dataset_names:
            ds_sub = ds.sel(time=slice(*config[f'{s}_date_range']))

            if s == 'train':
                for key, dropouts in dropoutConfig.items():
                    try:
                        ds_sub['PM25'].loc[:, dropouts[0], dropouts[1]] = np.nan
                        print(f"dropped {dropouts[0]}, {dropouts[1]}")
                    except KeyError:
                        print(f"{dropouts[0]}, {dropouts[1]} not in bounds")

            datasets[s] = EpaDataset(ds_sub, window_size, in_channels, out_channels)

    loaders = {
        k: DataLoader(v, batch_size=batch_size, shuffle=True)
        for k, v in datasets.items()
    }

    model = EpaSeq2Seq(
        in_channels=len(in_channels),
        out_channels=len(out_channels),
        frame_size=datasets['train'].frame_size,
        **model_params
    ).to(model_params['device'])
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    loss_function = mse_nan_loss
    if checkpoint_path and os.path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location=model_params['device'])
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    with mlflow.start_run():
        for ix_epoch in range(epochs + 1):

            if ix_epoch == 0:
                print("Untrained model\n--------")
                train_loss = eval_model(loaders['train'], model, loss_function)
            else:
                print(f"Epoch {ix_epoch}\n---------")
                train_loss = train_model(loaders['train'], model, loss_function, optimizer=optimizer)

            validation_loss = eval_model(loaders['validation'], model, loss_function)
            save_state(
                outputdir, config,
                ix_epoch, model, optimizer, train_loss, validation_loss,
                loss_function.__name__
            )
            mlflow.log_metrics({'train_loss': train_loss, 'validation_loss': validation_loss}, step=ix_epoch)
            print()


if __name__ == '__main__':
    main()