monkey_utils.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import auc
from sklearn.metrics import r2_score
from scipy.stats import ttest_rel

def train_network(net, train_loader, epochs = 50):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net.to(device)
    criterion = nn.MSELoss()
    optimizer = optim.SGD(net.parameters(), lr = 0.001, momentum=0.9)

    for epoch in range(epochs):
        
        running_loss = 0.0
        for i, data in enumerate(train_loader):
            inputs,labels = data

            inputs, labels = inputs.to(device), labels.to(device)
            if  np.isnan(inputs.cpu()).any().item() > 0 or  np.isnan(labels.cpu().any().item()) > 0:
                print(np.isnan(inputs.cpu()).any(), np.isnan(labels.cpu().any()))
            
            optimizer.zero_grad()
            outputs, states = net(inputs, reset_states = True)

            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if i % 500 == 499:  
                    print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 500:.3f}')
                    running_loss = 0.0
    return net

def test_network(net, test_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    r2_res = 0
    net.eval()
    with torch.no_grad():
        all_preds = []
        all_labels = []
        
        for data in test_loader:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            
            preds, states = net(inputs, reset_states = True)
            preds = preds.squeeze(); labels = labels.squeeze()
            all_preds.append(preds); all_labels.append(labels)
        
        all_preds, all_labels = torch.cat(all_preds).cpu(), torch.cat(all_labels).cpu()
        r2 = r2_score(all_labels, all_preds)
        return r2