nn_only_seq2seq.py

import os
import argparse
import logging
import sys

import time
import numpy as np
from src.utils import Params, save_checkpoint, load_checkpoint, set_logger, delete_file_or_folder
from src.utils import set_logger, delete_file_or_folder
from src.metrics import instantiate_losses, instantiate_metrics, functionalize_metrics

from src.dataset.idm_data import idm_data_loader

from src.models.neuralode_seq2seq import ODEFunc

import torch
import torch.nn as nn
import torch.optim as optim

from torchdiffeq import odeint_adjoint as odeint

parser = argparse.ArgumentParser('ODE demo')
parser.add_argument('--experiment_dir', default=None,
                    help="Directory containing experiment_setting.json")
parser.add_argument('--restore_from', default= None,
                    help="Optional, file location containing weights to reload")
parser.add_argument('--mode', default='train',
                    help="train, test, or train_and_test")
parser.add_argument('--force_overwrite', default=False, action='store_true',
                    help="For debug. Force to overwrite")
args = parser.parse_args()


device = torch.device('cpu')

# Set the random seed for the whole graph for reproductible experiments
if __name__ == "__main__":
    # Load the parameters from the experiment params.json file in model_dir
    args = parser.parse_args()
    json_path = os.path.join(args.experiment_dir, 'experiment_setting.json')
    assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)
    params = Params(json_path)

    if params.adjoint:
        from torchdiffeq import odeint_adjoint as odeint
    else:
        from torchdiffeq import odeint

    # if test, use cpu
    if args.mode == "test":
        device = torch.device('cpu')
        logging.info("In the test mode, use cpu")

    # check empty
    if os.path.exists(os.path.join(args.experiment_dir, 'train.log')) & args.force_overwrite is not True:
        raise Exception("previous experiment results exist")

    # Set the logger
    set_logger(os.path.join(args.experiment_dir, 'train.log'))
    logging.info(" ".join(sys.argv))

    # Create the input data pipeline
    logging.info("Loading the datasets...")

    data_loader = idm_data_loader(params.data)

    logging.info("Creating the model...")

    # create model
    input_dim = 2
    output_dim = 2
    net_args = (input_dim, output_dim,
              params.neural_ode["net"]["n_hidden"],
              params.neural_ode["net"]["hidden_dim"])

    net_kwargs = {"activation_type": params.neural_ode["net"]["activation_type"],
                "last_activation_type": params.neural_ode["net"]["last_activation_type"],
                "mean": data_loader.mean,
                "std": data_loader.std,
                "device": device}

    func = ODEFunc(device, net_args, net_kwargs).to(device)
    optimizer = optim.RMSprop(func.parameters(), lr=1e-3)

    for itr in range(1, 100 + 1):
        optimizer.zero_grad()
        batch_y0, batch_t, batch_y = data_loader.get_batch()
        pred_y = odeint(func, batch_y0, batch_t).to(device)
        loss = torch.mean(torch.abs(pred_y - batch_y))
        loss.backward()
        optimizer.step()

        # if itr % args.test_freq == 0:
        #     with torch.no_grad():
        #         pred_y = odeint(func, true_y0, t)
        #         loss = torch.mean(torch.abs(pred_y - true_y))
        #         print('Iter {:04d} | Total Loss {:.6f}'.format(itr, loss.item()))