Pytorch version of the codes

[Carbord]

Hello! I’m also trying to reproduce the code in pytorch. But my mse loss of validation set is about 1 at 150 epochs, where the loss is about 0.0065 when running your tensorflow codes. Will you share your pytorch version code? Or would you be kind to help checking the following reproduced codes? Thanks a lot!

[SizheAn]

Maybe try this, define a MARS_CNN_model.py:

# -*- coding: utf-8 -*-
"""
Created on Thu Nov  4 22:27:33 2021

@author: sizhean
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

torch.set_default_tensor_type(torch.FloatTensor)

class CNN(nn.Module):
    def __init__(self, channels, neurons):
        super(CNN, self).__init__()
        
        self.conv1 = nn.Conv2d(channels, 16, kernel_size=3, padding=1)
        self.drop1 = nn.Dropout(0.3)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.drop2 = nn.Dropout(0.3)
        self.batchnorm1 = nn.BatchNorm2d(32, momentum=0.05, eps = 0.001)
        
        self.fc = nn.Linear(2048, 512)
        self.batchnorm2 = nn.BatchNorm1d(512, momentum=0.05, eps = 0.001)
        self.drop3 = nn.Dropout(0.4)
        self.regression = nn.Linear(512, neurons)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.drop1(x)
        x = F.relu(x)
        
        x = self.conv2(x)
        x = self.drop2(x)
        x = F.relu(x)
        x = self.batchnorm1(x)
        
        # flatten
        x = x.reshape(x.size(0), -1)
        x = self.fc(x)
        x = F.relu(x)
        x = self.batchnorm2(x)
        x = self.drop3(x)
        
        x = self.regression(x)
        
        return x

In the main function:

from MARS_CNN_model import CNN
import torch.nn.functional as F
import torch.optim as optim

model = CNN(5,51).to(device)

creterion = F.l1_loss
optimizer = optim.Adam(model.parameters(), betas=(0.5, 0.999))

        # for the main training
        for epoch_count in range(total_epochs):
            ###################
            # train the model #
            ###################
            model.train()

            for batch_data, batch_labels in train_loader:
                batch_data = batch_data
                batch_data = batch_data.cuda()
                batch_labels = batch_labels.cuda()
                # forward
                batch_preds = model(batch_data)
                loss = creterion(batch_preds, batch_labels)
                # print("the current loss is: ", loss)
                # set the gradient to zero
                optimizer.zero_grad()
                # do backward prop
                loss.backward()
                # update the gradient
                optimizer.step()
                
                batch_iter += 1
                
                # record training loss
                train_losses.append(loss.item())
                # if batch_iter % 200 == 0:
                #     print('iter num: {} \t loss: {:.2f}'.format(batch_iter, loss.item()))

Note that one easy problem could be when you change the dimension from 885 to 588, make sure they are transforming the way you want. Sometime the dimension could be messed up.

Sorry for the inconvenience, I will try to upload a pytorch version when I got time.