Pytorch training example that can possibly deadlock

test_deadlock.py

import os
import sys
import shutil
import numpy as np
import torch
import torch.nn as nn
import random

import torchvision.transforms as transforms
from torchvision.models import alexnet
from torchvision.datasets import ImageFolder
from datetime import datetime

if __name__ == '__main__':
    torch.set_num_threads(40)
    torch.manual_seed(0)
    np.random.seed(0)
    random.seed(0)
    
    crop_size = 227
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomCrop(crop_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    dataset1 = ImageFolder('data/office_home/Real', transform)
    dataset2 = ImageFolder('data/office_home/Clipart', transform)

    loader1 = torch.utils.data.DataLoader(
        dataset1, num_workers=20, shuffle=True, 
        batch_size=32, drop_last=False, pin_memory=True)
    loader2 = torch.utils.data.DataLoader(
        dataset2, num_workers=20, shuffle=True, 
        batch_size=32, drop_last=False, pin_memory=True)
    
    loader3 = torch.utils.data.DataLoader(
        dataset1, num_workers=20, shuffle=True, 
        batch_size=32, drop_last=False, pin_memory=True)

    G = alexnet(pretrained=False, num_classes=65)
    G.train()
    G.cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(G.parameters(), lr=0.01)

    print('Starting training')
    steps = 100
    for step in range(steps):
        if step % 5 == 0:
            data_iter1 = iter(loader1)
        if step % 13 == 0:
            data_iter2 = iter(loader2)
        if step % 7 == 0:
            data_iter3 = iter(loader3)

        imgs1, labels1 = next(data_iter1)
        imgs2, labels2 = next(data_iter2)
        imgs3, labels3 = next(data_iter3)
        imgs1 = imgs1.cuda(non_blocking=True)
        imgs2 = imgs2.cuda(non_blocking=True)
        imgs3 = imgs3.cuda(non_blocking=True)
        labels1 = labels1.cuda(non_blocking=True)
        labels2 = labels2.cuda(non_blocking=True)
        labels3 = labels3.cuda(non_blocking=True)

        loss1 = criterion(G(imgs1), labels1)
        loss2 = criterion(G(imgs2), labels2)
        loss3 = criterion(G(imgs3), labels3)

        loss = loss1 + loss2 + loss3
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        now = str(datetime.now().strftime("%H:%M %d-%m-%Y"))
        print('[{}] Step [{}] Loss : {:.4f}'.format(now, step, loss.item()))