from torchvision import models, transforms
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

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

import numpy as np
import random


class BaseDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        assert len(images) == len(labels), 'image num must equal to label num!'
        self.images = images
        self.labels = labels
        self.transform = transform
        
    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        
        if self.transform:
            image = self.transform(image)
        return image, label
        
    def __len__(self):
        return len(self.images)


random.seed(2021)
np.random.seed(2021)
torch.manual_seed(2021)

# deterministic for multiprocess in dataloader
def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    print(f"seed of process-{ worker_id }: { worker_seed }")
    random.seed(worker_seed)
    np.random.seed(worker_seed)


epoch_num = 1000
batch_size = 32


model = models.resnet50()
if torch.cuda.is_available():
    model.cuda()


optimizer = optim.Adam(model.parameters(), lr=1e-4)


if __name__ == '__main__':
    print("---start training...")

    # torch.backends.cudnn.benchmark = True  # use this for faster training
    torch.backends.cudnn.deterministic = True  # use this for more deterministic training

    for epoch in range(1, epoch_num + 1):
        # reset loss and item
        running_loss = 0.0
        ite_num4loss = 0
        
        images = np.random.randn(1600, 224, 224, 3)
        labels = np.random.randint(1000, size=(1600))
        
        dataset = BaseDataset(images, labels, 
                            transform=transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize(
                                    mean=[0.484, 0.454, 0.403], 
                                    std=[0.225, 0.220, 0.220])]))
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0, 
                                worker_init_fn=seed_worker)  # set num_workers = 0 in jupyter


        model.train()

        for i, data in tqdm(enumerate(dataloader)):
            ite_num4loss = ite_num4loss + 1
    
            inputs, labels = data
            inputs = inputs.to(torch.float32)
            labels = labels.to(torch.int64)
            
            if torch.cuda.is_available():
                inputs, labels = inputs.cuda(), labels.cuda()
    
            # zero the parameter gradients
            optimizer.zero_grad()
    
            # forward + backward + optimize
            out = model(inputs)
            loss = F.cross_entropy(out, labels)
    
            loss.backward()
            optimizer.step()
    
            # print statistics
            running_loss += loss.item()
            
            print(f"[epoch: {epoch:3d}/{epoch_num:3d}" + "," +
                f" batch: {(i+1)*batch_size:6d}/{len(dataset)}" + "," +
                f" train loss: {running_loss / ite_num4loss:3f}")


I test the above code with torch1.8.2+cu102 and torch1.10.0+cu113 respectively. The results are as follows:


torch1.8.2+cu102, 2.73 it/s, 19 s/epoch


torch1.10.0+cu113, 2.36 it/s, 25 s/epoch


You can easily run this code in your command line.