Skip to content

Instantly share code, notes, and snippets.

@samarth4149
Last active Aug 26, 2020
Embed
What would you like to do?
Pytorch training example that can possibly deadlock
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()))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment