<jupyter_start><jupyter_text>Imports and Constants<jupyter_code># File system
import os, os.path 
import pickle
from glob import glob
import sys

# Workflow
import random
import tqdm
import tqdm.notebook

# Computation
import numpy as np
import torch
import torch.nn as nn

# Data visualization
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
sns.set(rc={"figure.dpi":100, 'savefig.dpi':100})
sns.set_context('notebook')

# Constants
from argoverse_constants import *

torch.__version__
# Keys to the pickle objects
CITY = 'city'
LANE = 'lane'
LANE_NORM = 'lane_norm'
SCENE_IDX = 'scene_idx'
AGENT_ID = 'agent_id'
P_IN = 'p_in'
V_IN = 'v_in'
P_OUT = 'p_out'
V_OUT = 'v_out'
CAR_MASK = 'car_mask'
TRACK_ID = 'track_id'

# Encode Miami as 0, Pittsburgh as 1
MIA = 'MIA'
PIT = 'PIT'  
CITY_MAP = {MIA: 0, PIT: 1}

# Transformed data keys
LANE_IN = 'closest_lanes_in'
NORM_IN = 'closest_lane_norms_in'
LANE_OUT = 'closest_lanes_out'
NORM_OUT = 'closest_lane_norms_out'
# Dataset variables
VALIDATION_PATH = './val_data/'
ORIGINAL_PATH = './original_train_data/'
TRANSFORMED_TRAIN_PATH = './transformed_train_data'
TRANSFORMED_VAL_PATH = './transformed_val_data'
STANDARDIZED_TRAIN_PATH = './standardized_train_data'
STANDARDIZED_VAL_PATH = './standardized_val_data'

# train_path = TRANSFORMED_TRAIN_PATH
# val_path = TRANSFORMED_VAL_PATH

train_path = STANDARDIZED_TRAIN_PATH
val_path = STANDARDIZED_VAL_PATH

# Path to model predictions on test set
PREDICTION_PATH = './my_submission.csv'
# Header of predictions CSV file
CSV_HEADER = ['ID,'] + ['v' + str(i) + ',' for i in range(1, 60)] + ['v60', '\n']

# Get list of all training file names
original_files = glob(os.path.join(train_path, '*'))

# 80-20 train-test split
train_files = random.sample(original_files, int(len(original_files) * 0.8))
test_files = list(set(original_files) - set(train_files))

# Validation
val_files = glob(os.path.join(val_path, '*'))
# Control variables
USE_ONE_AGENT = True
PREDICT_ONE_AGENT = True  # controls whether 60-agent inputs are mapped to just the target agent's outputs

NUM_AGENTS = 1 if USE_ONE_AGENT else 60
IN_LEN = 19
OUT_LEN = 30

# Controls how many features to use
N_FEAT_IN = 4
N_FEAT_OUT = 2

APPLY_NORM = True
APPLY_ST = False  # standardization
APPLY_MM = True  # min-max

# Batch variables
BATCH_SIZE_TRAIN = 64
BATCH_SIZE_TEST = BATCH_SIZE_TRAIN
BATCH_SIZE_VAL = 32
N_WORKERS = 4
# Windows doesn't support anything but N_WORKERS = 0 for the DataLoader
if 'win' in sys.platform:
    N_WORKERS = 0
N_WORKERS<jupyter_output><empty_output><jupyter_text>Dataset Loading and Batching<jupyter_code>def standardize(x, mean, std):
    return (x - mean) / std
def inv_standardize(x, mean, std):
    return (x * std) + mean

def min_max_scale(x, minval, maxval):
    return (x - minval) / (maxval - minval)
def inv_min_max_scale(x, minval, maxval):
    return (x * (maxval - minval)) + minval

def scale(x, minval, maxval, mean, std):
    if APPLY_MM:
        return min_max_scale(x, minval, maxval)
    else:
        return standardize(x, mean, std)

def descale(x, minval, maxval, mean, std):
    if APPLY_MM:
        return inv_min_max_scale(x, minval, maxval)
    else:
        return inv_standardize(x, mean, std)

def normalize_inpos(inp, city):
    if city == 0:
        inp[:, 0] = scale(inp[:, 0], MIA_PX_IN_MIN, MIA_PX_IN_MAX, MIA_PX_IN_MEAN, MIA_PX_IN_STD)
        inp[:, 1] = scale(inp[:, 0], MIA_PY_IN_MIN, MIA_PY_IN_MAX, MIA_PY_IN_MEAN, MIA_PY_IN_STD)
    else:
        inp[:, 0] = scale(inp[:, 0], PIT_PX_IN_MIN, PIT_PX_IN_MAX, PIT_PX_IN_MEAN, PIT_PX_IN_STD)
        inp[:, 1] = scale(inp[:, 0], PIT_PY_IN_MIN, PIT_PY_IN_MAX, PIT_PY_IN_MEAN, PIT_PY_IN_STD)        
    return inp

def normalize_invel(inp, city):
    if city == 0:
        inp[:, 0] = scale(inp[:, 0], MIA_VX_IN_MIN, MIA_VX_IN_MAX, MIA_VX_IN_MEAN, MIA_VX_IN_STD)
        inp[:, 1] = scale(inp[:, 0], MIA_VY_IN_MIN, MIA_VY_IN_MAX, MIA_VY_IN_MEAN, MIA_VY_IN_STD)
    else:
        inp[:, 0] = scale(inp[:, 0], PIT_VX_IN_MIN, PIT_VX_IN_MAX, PIT_VX_IN_MEAN, PIT_VX_IN_STD)
        inp[:, 1] = scale(inp[:, 0], PIT_VY_IN_MIN, PIT_VY_IN_MAX, PIT_VY_IN_MEAN, PIT_VY_IN_STD)        
    return inp 


def normalize_outpos(inp, city):
    if city == 0:
        inp[:, 0] = scale(inp[:, 0], MIA_PX_OUT_MIN, MIA_PX_OUT_MAX, MIA_PX_OUT_MEAN, MIA_PX_OUT_STD)
        inp[:, 1] = scale(inp[:, 0], MIA_PY_OUT_MIN, MIA_PY_OUT_MAX, MIA_PY_OUT_MEAN, MIA_PX_OUT_STD)
    else:
        inp[:, 0] = scale(inp[:, 0], PIT_PX_OUT_MIN, PIT_PX_OUT_MAX, PIT_PX_OUT_MEAN, PIT_PX_OUT_STD)
        inp[:, 1] = scale(inp[:, 0], PIT_PY_OUT_MIN, PIT_PY_OUT_MAX, PIT_PY_OUT_MEAN, PIT_PY_OUT_STD)        
    return inp 


def normalize_outvel(inp, city):
    if city == 0:
        inp[:, 0] = scale(inp[:, 0], MIA_VX_OUT_MIN, MIA_VX_OUT_MAX, MIA_VX_OUT_MEAN, MIA_VX_OUT_STD)
        inp[:, 1] = scale(inp[:, 0], MIA_VY_OUT_MIN, MIA_VY_OUT_MAX, MIA_VY_OUT_MEAN, MIA_VX_OUT_STD)
    else:
        inp[:, 0] = scale(inp[:, 0], PIT_VX_OUT_MIN, PIT_VX_OUT_MAX, PIT_VX_OUT_MEAN, PIT_VX_OUT_STD)
        inp[:, 1] = scale(inp[:, 0], PIT_VY_OUT_MIN, PIT_VY_OUT_MAX, PIT_VY_OUT_MEAN, PIT_VY_OUT_STD)        
    return inp  


def inv_norm_outpos(out, cities):
    batch_size = out.shape[0]
    for b in range(batch_size):
        if cities[b] == 0:  # Miami
            out[b, :, 0] = descale(out[b, :, 0], MIA_PX_OUT_MIN, MIA_PX_OUT_MAX, MIA_PX_OUT_MEAN, MIA_PX_OUT_STD)
            out[b, :, 1] = descale(out[b, :, 1], MIA_PY_OUT_MIN, MIA_PY_OUT_MAX, MIA_PY_OUT_MEAN, MIA_PX_OUT_STD)
        else:  # Pittsburgh
            out[b, :, 0] = descale(out[b, :, 0], PIT_PX_OUT_MIN, PIT_PX_OUT_MAX, PIT_PX_OUT_MEAN, PIT_PX_OUT_STD)
            out[b, :, 1] = descale(out[b, :, 1], PIT_PY_OUT_MIN, PIT_PY_OUT_MAX, PIT_PY_OUT_MEAN, PIT_PX_OUT_STD) 
    return out
class ArgoverseDataset(torch.utils.data.Dataset):
    def __init__(self, files):
        super(ArgoverseDataset, self).__init__()
        self.files = files
        
    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):

        pkl_path = self.files[idx]
        with open(pkl_path, 'rb') as f:
            data = pickle.load(f)
        return data
def collate_train_test(batch):
    """ 
    Custom collate_fn function to be used for DataLoader.    
    The input tensor is organized as 19 rows, where each row has 4 columns: px, py, vx, vy
    
    The output tensor can be organized in 2 ways:
    A) The output tensor is organized as 120 rows, where each row has 1 column. 
    Each contiguous sequence of 4 elements is px, py, vx, vy
    
    B) The input tensor is organized as 30 rows, where each row has 4 columns: px, py, vx, vy
    """    
    inp = []        
    out = []
    agent_idxs = []
    scene_idxs = []
    city_ids = []
    for scene in batch:
        # Get the target agent id
        agent_id = scene[AGENT_ID]        
        # Get the matrix of all agents
        track_id = scene[TRACK_ID]        
        # Get the location of the target agent in the matrix
        idx = np.nonzero(track_id[:, 0] == agent_id)[0][0]
        
        # Number of time steps in the input and output sequences
        inlen = scene[P_IN].shape[1]
        outlen = scene[P_OUT].shape[1]
        
        # Aliases of scene variables for convenience
        pin, pout, vin, vout = scene[P_IN], scene[P_OUT], scene[V_IN], scene[V_OUT]
        lanein, laneout, normin, normout = scene[LANE_IN], scene[LANE_OUT], scene[NORM_IN], scene[NORM_OUT]
        
        # Use these lines to include only the target agent.
        # Training with just one agent is faster and preferred for experimentation
        if USE_ONE_AGENT:  
            if scene[CITY] == MIA:
                cities = np.zeros((inlen, 1)) 
            else:
                cities = np.ones((inlen, 1)) 

            # Use this line to include all input features
#             inp_tens = np.concatenate((pin[idx], vin[idx], lanein[idx], normin[idx], cities), axis=1)    

            # Lane only
#             inp_tens = np.concatenate((pin[idx], vin[idx], lanein[idx], normin[idx]), axis=1) 

            # Only include cities
#             inp_tens = np.concatenate((pin[idx], vin[idx], cities), axis=1)     

            # Only include pos/vel
#             if APPLY_NORM:
#                 pin[idx] = normalize_inpos(pin[idx], CITY_MAP[scene[CITY]])
#                 vin[idx] = normalize_invel(vin[idx], CITY_MAP[scene[CITY]])
#                 pout[idx] = normalize_outpos(pout[idx], CITY_MAP[scene[CITY]])

            inp_tens = np.concatenate((pin[idx], vin[idx]), axis=1)

            # All output features
#             out_tens = np.concatenate((pout[idx], vout[idx], laneout[idx], normout[idx]), axis=1)
            # Only include output position and velocity
#             out_tens = np.concatenate((pout[idx], vout[idx]), axis=1)   
        
            out_tens = pout[idx]
            
        # Use these lines to include all 60 agents, even the dummy agents.
        # Training with all 60 agents is much slower than with just one agent.
        else:
            num_agents = pin.shape[0]
            if scene[CITY] == MIA:
                cities = np.zeros((num_agents, inlen, 1)) 
            else:
                cities = np.ones((num_agents, inlen, 1))  
                
            # Include all features
#             inp_tens = np.concatenate((pin, vin, lanein, normin, cities), axis=2)  
            # Lane only
#             inp_tens = np.concatenate((pin, vin, lanein, normin), axis=2)   
            # Include only cities
            inp_tens = np.concatenate((pin, vin, cities), axis=2) 
            # Include only pos/vel
#             inp_tens = np.concatenate((pin, vin), axis=2)
            
            # Makes predictions for only the target agent
            if PREDICT_ONE_AGENT:
                # All features
#                 out_tens = np.concatenate((pout[idx], vout[idx], laneout[idx], normout[idx]), axis=1)                
                # Pos/vel
                out_tens = np.concatenate((pout[idx], vout[idx]), axis=1)

            # Makes predictions for all 60 agents
            else:
                # All features
#                 out_tens = np.concatenate((pout, vout, laneout, normout), axis=2)                
                # pv
                out_tens = np.concatenate((pout, vout), axis=2) 
        
        inp.append(inp_tens)
        out.append(out_tens)        
        scene_idxs.append(scene[SCENE_IDX]) 
        agent_idxs.append(idx)
        city_ids.append(CITY_MAP[scene[CITY]])

    inp = torch.FloatTensor(inp)
    out = torch.FloatTensor(out)
    return [inp, out, scene_idxs, agent_idxs, city_ids]
def collate_val(batch):
    """ 
    Custom collate_fn for validation dataset. The validation data do not contain output values.   
    The input tensor is organized as 19 rows, where each row has 4 columns: px, py, vx, vy
    """   
    inp = []
    scene_idxs = []
    agent_idxs = []
    city_ids = []
    
    for scene in batch:
        # Get the target agent id
        agent_id = scene[AGENT_ID]        
        # Get the matrix of all agents
        track_id = scene[TRACK_ID]        
        # Get the location of the target agent in the matrix
        idx = np.nonzero(track_id[:, 0] == agent_id)[0][0]
        
        inlen = scene[P_IN].shape[1]
        
        # Aliases of scene variables for convenience
        pin, vin = scene[P_IN], scene[V_IN]
        lanein, normin = scene[LANE_IN], scene[NORM_IN]
        num_agents = scene[P_IN].shape[0]
        
        # Use these lines to include only the target agent.
        if USE_ONE_AGENT:
            if scene[CITY] == MIA:
                cities = np.zeros((inlen, 1)) 
            else:
                cities = np.ones((inlen, 1)) 
            # All
#             inp_tens = np.concatenate((pin[idx], vin[idx], lanein[idx], normin[idx], cities), axis=1)
            # Lane only
#             inp_tens = np.concatenate((pin[idx], vin[idx], lanein[idx], normin[idx]), axis=1)
            # Cities
#             inp_tens = np.concatenate((pin[idx], vin[idx], cities), axis=1)  
            # Pos/vel
#             if APPLY_NORM:
#                 pin[idx] = normalize_inpos(pin[idx], CITY_MAP[scene[CITY]])
#                 vin[idx] = normalize_invel(vin[idx], CITY_MAP[scene[CITY]])
            inp_tens = np.concatenate((pin[idx], vin[idx]), axis=1)  

        # Use these lines to include all 60 agents
        else:
            if scene[CITY] == MIA:
                cities = np.zeros((num_agents, inlen, 1)) 
            else:
                cities = np.ones((num_agents, inlen, 1))
            # All features
#             inp_tens = np.concatenate( (pin, vin, lanein, normin, cities), axis=2)       
            # Lane only
#             inp_tens = np.concatenate( (pin, vin, lanein, normin), axis=2)
            # Cities
            inp_tens = np.concatenate( (pin, vin, cities), axis=2)             
            # pv
#             inp_tens = np.concatenate( (pin, vin), axis=2)  
        
        inp.append(inp_tens)        
        scene_idxs.append(scene[SCENE_IDX])
        agent_idxs.append(idx)
        city_ids.append(CITY_MAP[scene[CITY]])        
        
    inp = torch.FloatTensor(inp)    
    return [inp, scene_idxs, agent_idxs, city_ids]
# Initiliaze datasets and loaders

train_dataset = ArgoverseDataset(train_files)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE_TRAIN, 
                                           shuffle=True, collate_fn=collate_train_test, 
                                           num_workers=N_WORKERS)
test_dataset = ArgoverseDataset(test_files)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE_TEST, 
                                          shuffle=False, collate_fn=collate_train_test, 
                                          num_workers=N_WORKERS)
val_dataset = ArgoverseDataset(val_files)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE_VAL, 
                                         shuffle=False, collate_fn=collate_val,
                                         num_workers=N_WORKERS)
# Look at one data sample      
for _, (data, target, agent_idxs, scene_idxs, cities) in enumerate(train_loader):
    print(data.shape)
    print(target.shape) 
        
    break
# Look at one data sample      
for _, (data, agent_idxs, scene_idxs, masks) in enumerate(val_loader):
    print(data.shape)
    break<jupyter_output>torch.Size([32, 19, 4])<jupyter_text>Training Workflow<jupyter_code>def train(model, device, train_loader, optimizer, epoch):
    # Set the model into training mode
    model.train()    
    
    # Define the loss function.
    criterion = torch.nn.MSELoss(reduction='mean')
    total_loss = 0    
    for i in range(epoch):
        h = c = None # Hidden states      

        iterator = tqdm.notebook.tqdm(train_loader, total=int(len(train_loader)))
        for _, batch in enumerate(iterator):
            data, target, scene_idxs, agent_idxs, cities = batch
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()            
    
            out = model(data)
            # Compute the loss
            # Rescale values
            out = inv_norm_outpos(out, cities)
            target = inv_norm_outpos(target, cities)
            
            loss = torch.sqrt(criterion(out, target))            
            total_loss += loss.item()

            # Perform backpropagation
            loss.backward()
            # Update the weights
            optimizer.step()

            # Update the progress bar for tqdm
            iterator.set_postfix(train_loss=loss.item())
            
    return (total_loss * BATCH_SIZE_TRAIN) / len(train_files)
def test(model, device, test_loader):
    model.eval()    
    criterion = torch.nn.MSELoss(reduction='mean')    
    iterator = tqdm.notebook.tqdm(test_loader, total=int(len(test_loader)))
    total_loss = 0
    
    for _, batch in enumerate(iterator):
        data, target, scene_idxs, agent_idxs, masks = batch
        data, target = data.to(device), target.to(device)
        
        with torch.no_grad():
            
            out = model(data)
            # Rescale values
            out = inv_norm_outpos(out, cities)
            target = inv_norm_outpos(target, cities)
            
            loss = torch.sqrt(criterion(out, target))        
            total_loss += loss.item()
                
            iterator.set_postfix(test_loss=loss.item())
        
    return (total_loss * BATCH_SIZE_TEST) / len(test_files)
def train_test(model, DEVICE, train_loader, optimizer, NUM_EPOCH):
    for t in range(1, NUM_EPOCH + 1):
        train_loss = train(model, DEVICE, train_loader, optimizer, 1)
        test_loss = test(model, DEVICE, test_loader)
        
        train_losses.append(train_loss)
        test_losses.append(test_loss)
        
        print(f'Epoch {t}: train_loss = {train_loss}, test_loss = {test_loss}')         
def validate(model, device, val_loader, path):
    """
    path: path to csv file to write predictions
    """
    model.eval()   
    
    # Prep the output file
    with open(PREDICTION_PATH, "w") as csv_file:
        # Clear the csv file before appending data to it
        csv_file.truncate()
        # Write the header to the csv file
        csv_file.writelines(CSV_HEADER)    
    
    # Make predictions
    with open(path, "a") as pred_file:        
        iterator = tqdm.notebook.tqdm(val_loader, total=int(len(val_loader)))
        
        for _, batch in enumerate(iterator):
            data, scene_idxs, agent_idxs, cities = batch
            data = data.to(device) 
            
            with torch.no_grad():
                output = model(data)
                output = inv_norm_outpos(data, cities)
                # Convert the Tensor from GPU -> CPU -> NumPy array
                np_out = output.cpu().detach().numpy()
                
                # Store only the predictions for the target agent and keep the positions, not the velocities
                batch_size = np_out.shape[0]
                
                pred = np.zeros((batch_size, 60))
                if PREDICT_ONE_AGENT:
                    # The output should be a (batch size, time steps, num features out) tensor 
                    # where the first two features are the out position x, y
                    for i in range(batch_size):
                        pred[i] = np_out[i, :, :2].flatten()
                else:
                    # The output should be a (batch size, num agents, time steps, num features out) tensor 
                    # where the first two features are the input position and the output position
                    for i in range(batch_size):
                        idx = agent_idxs[i]
                        pred[i] = np_out[i, idx, :, :2].flatten()                    

                # Form comma-separated string
                s = []
                for i in range(pred.shape[0]):
                    s.append(','.join([str(scene_idxs[i])] + [str(v) for v in pred[i]]) + '\n')

                # Write data to file
                pred_file.writelines(s)<jupyter_output><empty_output><jupyter_text>Model Initialization<jupyter_code>class ArgoNet(torch.nn.Module):
    """
    Neural Network class - linear regression
    """
    def __init__(self, device):
        super(ArgoNet, self).__init__() 
        
        self.device = device        
        # Linear regression for 1 agent     
        if PREDICT_ONE_AGENT:
            self.fc = nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, 1 * OUT_LEN * N_FEAT_OUT)
        else:
            self.fc = nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, NUM_AGENTS * OUT_LEN * N_FEAT_OUT)
    
    def forward(self, x):
        x = x.view(x.size(0), -1)  # flatten result
        x = self.fc(x)            
        if PREDICT_ONE_AGENT:
            x = x.view(x.size(0), OUT_LEN, N_FEAT_OUT)
        else:
            x = x.view(x.size(0), NUM_AGENTS, OUT_LEN, N_FEAT_OUT)
        return x
# class ArgoNet(torch.nn.Module):
#     """
#     Neural Network class - 1 hidden layer
#     """
#     def __init__(self, device):
#         super(ArgoNet, self).__init__() 
        
#         self.device = device 
#         self.hidden_size = IN_LEN * N_FEAT_IN * 30
#         if PREDICT_ONE_AGENT:
#             self.fc = nn.Sequential(
#                 nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, 1 * self.hidden_size),
#                 nn.ReLU(),
#                 nn.Linear(self.hidden_size, 1 * OUT_LEN * N_FEAT_OUT)
#             )
#         else:
#             self.fc = nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, NUM_AGENTS * OUT_LEN * N_FEAT_OUT)
    
#     def forward(self, x):
#         x = x.view(x.size(0), -1)  # flatten result
#         x = self.fc(x)            
#         if PREDICT_ONE_AGENT:
#             x = x.view(x.size(0), OUT_LEN, N_FEAT_OUT)
#         else:
#             x = x.view(x.size(0), NUM_AGENTS, OUT_LEN, N_FEAT_OUT)
#         return x
# class ArgoNet(torch.nn.Module):
#     """
#     Neural Network class - 2 hidden layers
#     """
#     def __init__(self, device):
#         super(ArgoNet, self).__init__() 
        
#         self.device = device 
#         # Linear regression for 1 agent     
#         if PREDICT_ONE_AGENT:
#             self.fc = nn.Sequential(
#                 nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, 1 * 512),
#                 nn.ReLU(),
#                 nn.Linear(512, 1024),
#                 nn.ReLU(),
#                 nn.Linear(1024, 1 * OUT_LEN * N_FEAT_OUT)
#             )
#         else:
#             self.fc = nn.Linear(NUM_AGENTS * IN_LEN * N_FEAT_IN, NUM_AGENTS * OUT_LEN * N_FEAT_OUT)
    
#     def forward(self, x):
#         x = x.view(x.size(0), -1)  # flatten result
#         x = self.fc(x)            
#         if PREDICT_ONE_AGENT:
#             x = x.view(x.size(0), OUT_LEN, N_FEAT_OUT)
#         else:
#             x = x.view(x.size(0), NUM_AGENTS, OUT_LEN, N_FEAT_OUT)
#         return x
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
DEVICE
# Code to save and reload a model
MODEL_STATE = 'model_state_dict'
OPTIMIZER_STATE = 'optimizer_state_dict'
EPOCH_STATE = 'epoch'
LOSS_STATE = 'loss'
BATCH_STATE = 'batch'

def save_model(path, model_state_dict, optimizer_state_dict, epoch, loss, batch):
    to_save = {
        MODEL_STATE: model_state_dict,
        OPTIMIZER_STATE: optimizer_state_dict,
        EPOCH_STATE: epoch,
        LOSS_STATE: loss,
        BATCH_STATE: batch
    }
    torch.save(to_save, path)
    
def load_model(path, model_to_load, optimizer_to_load):
    checkpoint = torch.load(path)
    model_to_load.load_state_dict(checkpoint[MODEL_STATE])
    optimizer_to_load.load_state_dict(checkpoint[OPTIMIZER_STATE])
    return checkpoint[EPOCH_STATE], checkpoint[LOSS_STATE], checkpoint[BATCH_STATE]
model = ArgoNet(DEVICE).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters())
# learning_rate = 1e-1
# momentum = 0.9
# weight_decay = 1
# optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
# optimizer = torch.optim.Adadelta(model.parameters())  # Used to generate visualizations in report
# optimizer = torch.optim.AdamW(model.parameters(), betas=(0.5, 0.9))

num_params = sum(p.numel() for p in model.parameters())   
print(f"Number of model parameters is {num_params}")
NUM_EPOCH = 3

# used for visualizing the loss
train_losses = []
test_losses = []
# Use these two lines to save a model
# save_model('two_hid_one_agent_ep6.tar', model.state_dict(), optimizer.state_dict(), NUM_EPOCH, 
#            (train_losses[-1], test_losses[-1]), BATCH_SIZE_TRAIN)
# Reload a model
# model = ArgoNet(DEVICE).to(DEVICE)
# optimizer = torch.optim.Adam(model.parameters())
# load_model('one_hid_one_agent.tar', model, optimizer)
# train_losses.clear()
# test_losses.clear()
# num_params = sum(p.numel() for p in model.parameters())   
# print(f"Number of model parameters is {num_params}")<jupyter_output><empty_output><jupyter_text>Evaluation<jupyter_code>train_test(model, DEVICE, train_loader, optimizer, 1)
validate(model, DEVICE, val_loader, PREDICTION_PATH)<jupyter_output><empty_output><jupyter_text>Loss Visualization<jupyter_code>def visualize_loss(losses):
    """
    Plots the losses over each training iteration. 
    Assumes that each element of the 'losses' list corresponds to the loss after each batch of train()
    """
    t_iter = np.arange(1, len(losses) + 1, 1, dtype=int)
    ax = sns.scatterplot(x=t_iter, y=losses, alpha=0.5)    
    ax.set_xlabel('Batch iteration number')
    ax.set_ylabel('Root-mean-square loss')
    ax.set_title('Batch Iteration vs. Root-Mean-Square Loss')
    plt.savefig('lossViter')
visualize_loss(train_losses)<jupyter_output><empty_output><jupyter_text>Ground Truth Comparison<jupyter_code>def visualize_predictions(model, device, loader):
    """
    Compares some randomly selected data samples to the model's predictions
    """
    model.eval()
    
    # Get a batch of data
    _, (inp, out, scene_idxs, agent_idxs, masks) = next(enumerate(loader))
    
    # Move tensors to chosen device
    inp, out = inp.to(device), out.to(device)
    
    # Sample number
    i = 0
    
    # Scene idx
    scene_idx = scene_idxs[i]
        
    # Get contiguous arrays of the ground truth output positions
    truth = target[i].cpu().detach().numpy()
    x = truth[:, 0]
    y = truth[:, 1] 
        
    # Get contiguous arrays of the prediction output positions
    output = model(inp)    
    pred = output[i].cpu().detach().numpy()
    xh = pred[:, 0]
    yh = pred[:, 0]    
    
    # Plot the ground truth and prediction positions
    fig, (ax) = plt.subplots(nrows=1, ncols=1, figsize=(3, 3))
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title('Scene ' + str(scene_idx))
    ax.scatter(x, y, label='Ground Truth')
    ax.scatter(xh, yh, label='Prediction')
    ax.legend()
visualize_predictions(model, DEVICE, train_loader)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11013417944111013, 'fraction_numerical': 0.008792385030008792, 'mean_word_length': 2.3864077669902914, 'pattern_counts': {'":': 1, '<': 28, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '78467', 'license': 'MIT', 'n_tokens_mistral': 9371, 'n_tokens_neox': 9031, 'n_words': 2072, 'path': 'previous/argoverse_norm.ipynb', 'repo_name': 'and-le/cse-151b-argoverse', 'size': '26159'}