[PYTHON] PyTorch learning template

Introduction

A description template for PyTorch datasets, models, and learning. It's almost for myself, so it may not be helpful.

version

Python 3.7.9 torch 1.6.0 torchvision 0.7.0

Data set creation

Creating your own dataset

load_data.py

import torch

def load_data(file_path):
    '''
Data read
    '''
    return input_data, output_data

class myDataset(torch.utils.data.Dataset):
    def __init__(self, file_path):
        input_data, output_data = load_data(file_path)
        self.input = torch.tensor(input_data, dtype=torch.float)
        self.data_num = len(self.input)
        self.output = torch.tensor(output_data, dtype=torch.float)

    def __len__(self):
        return self.data_num

    def __getitem__(self, idx):
        input = self.input[idx]
        output = self.output[idx]

        return input, output

model

Neural network model description

model.py

import torch
import torch.nn as nn

class myModel(nn.Module):
    def __init__(self, input_dim=10, fc_dim=300, out_dim=5):
        super().__init__()

        self.fc1 = nn.Linear(input_dim, fc_dim)
        self.fc2 = nn.Linear(fc_dim, out_dim)

        self.act = nn.ReLU()

    def forward(self, input):
        out = self.fc1(input)
        out = self.act(out)
        out = self.fc2(out)

        return out

Learning

train.py

import torch
import random
from model import myModel
from load_data import myDataset

torch.manual_seed(12)
random.seed(12)

#########Hyperparameters##########
#Model parameters
input_dim = 16
fc_dim = 1024
out_dim = 10

#Parameters during learning
batch_size = 8
learning_rate = 0.001
num_epochs = 100

#Verification interval
val_step = 5

#########Preparation for learning##########
#Create dataloader
file_path = './dataset'
dataset = myDataset(file_path)

#20% of the dataset is used as verification data
n_samples = len(dataset)
index_list = list(range(n_samples))
random.shuffle(index_list)
train_index = index_list[:int(n_samples * 0.8)]
val_index = index_list[int(n_samples * 0.8):]

train_dataset = torch.utils.data.dataset.Subset(dataset, train_index)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size, shuffle=True)

val_dataset = torch.utils.data.dataset.Subset(dataset, val_index)
val_dataloader = torch.utils.data.DataLoader(
    val_dataset, batch_size, shuffle=False)

dataloaders = {'train': train_dataloader, 'val': val_dataloader}

#GPU settings
device = torch.device(
    'cuda:0' if torch.cuda.is_available() else "cpu")

#Model settings
net = myModel(input_dim=input_dim, fc_dim=fc_dim, out_dim=out_dim)
net.to(device)

#Optimization method setting
# optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9)
optimizer = optim.Adam(net.parameters(), lr=learning_rate)

#Loss function settings
# criterion = nn.MSELoss()
criterion = nn.CrossEntropyLoss()

#########Learning##########
loss_list = {'train': [], 'val': []}

for epoch, phase in itertools.product(range(num_epochs), ['train', 'val']):
    if phase == 'train':
        #Training mode
        net.train()
    elif (epoch + 1) % val_step == 0:
        #Verification mode
        net.eval()
    else:
        continue

    #epoch loss sum
    epoch_loss = 0.0

    for inputs, results in tqdm(dataloaders[phase]):
        inputs_gpu = inputs.to(device)
        results_gpu = results.to(device)

        #optimizer initialization
        optimizer.zero_grad()

        #Forward propagation
        with torch.set_grad_enabled(phase == 'train'):
            outputs = net(inputs_gpu)

            loss = criterion(outputs, results_gpu)

            #Backpropagation
            if phase == 'train':
                loss.backward()
                optimizer.step()

            epoch_loss += loss.item() * inputs[1].size(0)

    epoch_loss = epoch_loss / len(dataloaders[phase].dataset)
    loss_list[phase].append(epoch_loss)

    #View results
    print('Epoch {}/{} | {:^5} | Loss: {}'.format(epoch +
                                                  1, num_epochs, phase, epoch_loss))