<jupyter_start><jupyter_text>Library<jupyter_code>import numpy as np
import torch
import torch.nn as nn

from utils import *
from dataset import CollisionDataset
from torch.utils.data import DataLoader<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>class ResidualMLP(nn.Module):

    def __init__(self, in_dim, out_dim):
        super(ResidualMLP, self).__init__()
        self.hidden_dim = 128
        self.fc_1 = nn.Sequential(
            nn.Linear(in_dim, self.hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            nn.ReLU(inplace=True),
        )
        self.fc_out = nn.Linear(self.hidden_dim, out_dim)

    def forward(self, x, y_phy):
        x = self.fc_1(x)
        x = self.fc_out(x)

        return x + y_phy
def train_model(model, train_loader, test_loader, num_epochs, optimizer, scheduler, criterion):
    model.train()
        
    # Training the Model
    min_test_dif = float('inf')
    epoch_loss = []
    for epoch in range(num_epochs):
        batch_loss = []
        for i, data in enumerate(train_loader):
            
            # get the inputs
            x = data['x']
            y = data['y']
            y_phy = data['y_phy']

            x = x.cuda()
            y = y.cuda()
            y_phy = y_phy.cuda()

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            y_hat = net(x, y_phy)
            
            loss = criterion(y_hat, y)
            
            loss.backward()
            optimizer.step()
            
            batch_loss.append(loss.item())

        # Results every epoch
        cur_epoch_loss = sum(batch_loss) / len(batch_loss)
        
        # Scheduler
        scheduler.step(cur_epoch_loss)
        
        # Test the network
        train_dif = test_model(model, train_loader, criterion)
        test_dif = test_model(model, test_loader, criterion)
        
        # Print the result
        print('Epoch: %d Train Loss: %f Train Dif: %f Test Dif: %f' 
              % (epoch, cur_epoch_loss, train_dif, test_dif))
        epoch_loss.append(cur_epoch_loss)
        
        if min_test_dif > test_dif:
            min_test_dif = test_dif
            print('Best')
        
    return epoch_loss

def test_model(model, test_loader, criterion):
    # Test the Model
    model.eval()
    
    batch_loss = []
    for i, data in enumerate(test_loader):

        # get the inputs
        x = data['x']
        y = data['y']
        y_phy = data['y_phy']

        x = x.cuda()
        y = y.cuda()
        y_phy = y_phy.cuda()

        y_hat = net(x, y_phy)
        loss = criterion(y_hat, y)
        batch_loss.append(loss.item())

    # Results every epoch
    cur_epoch_loss = sum(batch_loss) / len(batch_loss)
    
    model.train()
    
    return cur_epoch_loss
#################### Hyperparameters ####################
num_epochs = 50000
learning_rate = 0.001
weight_decay = 0
in_frames_num = 3
pre_frames_num = 15
factor = 0.95
patience = 40
batch_size = 16
#################### Hyperparameters ####################
net = ResidualMLP(in_dim=7, out_dim=2).cuda()
criterion = torch.nn.MSELoss()

train_set = CollisionDataset(
    './dataset/uIsPoint3/train', 
    sample_num=32
)

test_set = CollisionDataset(
    './dataset/uIsPoint3/test',
)

print(len(train_set), len(test_set))

train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_set, batch_size=len(test_set), shuffle=False)

optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=weight_decay)

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, 
    mode='min', 
    factor=factor, 
    patience=patience, 
    verbose=True, 
    threshold=1e-3
)

train_loss = train_model(
    net, 
    train_loader, 
    test_loader, 
    num_epochs, 
    optimizer, 
    scheduler, 
    criterion
)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12046080641121963, 'fraction_numerical': 0.008264462809917356, 'mean_word_length': 2.3422594142259414, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '64319', 'license': 'MIT', 'n_tokens_mistral': 1388, 'n_tokens_neox': 1335, 'n_words': 262, 'path': 'collision/Residual_physics_collision.ipynb', 'repo_name': 'physicslearning/PhysicsNAS', 'size': '3993'}