I have a very simple code which is supposed to run on one gpu. However, whenever I run a program on one gpu and trigger another run on another gpu, both of the programs either get stuck or slow down significantly.
from __future__ import print_function
from __future__ import division
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import time
import os
import copy
import argparse
#print("PyTorch Version: ",torch.__version__)
#print("Torchvision Version: ",torchvision.__version__)
# Number of classes in the dataset
num_classes = 2 # for now
# Batch size for training (change depending on how much memory you have)
batch_size = 32
# Number of epochs to train for
num_epochs = 50
# Flag for feature extracting. When False, we finetune the whole model,
# when True we only update the reshaped layer params
feature_extract = True
unfreeze_layer = 0
# Implement early stop
early_stop_tol = 5
parser = argparse.ArgumentParser()
parser.add_argument("gpu", type=str,
help="gpu number (e.g. \"0\" )")
parser.add_argument("unfreeze_layer", type=int)
parser.add_argument("path", type=str, help="path to input dir")
parser.add_argument("base_path", type=str, help="path to this python dir")
args = parser.parse_args()
data_dir = args.path
def unfreeze(model, min_layer):
print("params to learn:")
ct = 0
for child in model_ft.children():
ct += 1
if ct >= min_layer:
for name,param in model_ft.named_parameters():
param.requires_grad = True
print("\t",name)
def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, init_acc = 0.0):
since = time.time()
val_acc_history = []
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = init_acc
tol = 0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
model.train() # Set model to training mode
else:
model.eval() # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(dataloaders[phase].dataset)
epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
if phase == 'train':
writer.add_scalar("Loss/train", epoch_loss, epoch)
writer.add_scalar("Acc/train", epoch_acc, epoch)
if phase == 'val':
writer.add_scalar("Loss/val", epoch_loss, epoch)
writer.add_scalar("Acc/val", epoch_acc, epoch)
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
# save model
# torch.save(model, os.path.join(model_output_path, "model_" + str(epoch)))
# deep copy the model
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
if phase == 'val':
val_acc_history.append(epoch_acc)
tol += 1
if tol >= early_stop_tol:
print("reached max tol")
break
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(best_model_wts)
return model, val_acc_history
def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False
def initialize_model(num_classes, feature_extract, use_pretrained=True):
"""
Initialize input model.
"""
model_ft = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
input_size = 224
return model_ft, input_size
for learning_rate in np.geomspace(0.000001, 0.1, 100):
for finetune_learning_rate in np.geomspace(0.0000001, 0.01, 80):
# Initialize the model for this run
model_ft, input_size = initialize_model(num_classes, feature_extract, use_pretrained=True)
# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
'train': transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
# Create training and validation datasets
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']}
# Create training and validation dataloaders
dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=0) for x in ['train', 'val']}
images, labels = next(iter(dataloaders_dict['train']))
# Detect if we have a GPU available
device = torch.device(args.gpu if torch.cuda.is_available() else "cpu")
# Send the model to GPU
model_ft = model_ft.to(device)
# Gather the parameters to be optimized/updated in this run. If we are
# finetuning we will be updating all parameters. However, if we are
# doing feature extract method, we will only update the parameters
# that we have just initialized, i.e. the parameters with requires_grad
# is True.
params_to_update = model_ft.parameters()
print("Params to learn:")
if feature_extract:
params_to_update = []
for name,param in model_ft.named_parameters():
if param.requires_grad == True:
params_to_update.append(param)
print("\t",name)
else:
for name,param in model_ft.named_parameters():
if param.requires_grad == True:
print("\t",name)
writer = SummaryWriter(os.path.join(os.path.join(args.base_path, 'runs'), "batch_size" + str(learning_rate) + "_lr_" + str(finetune_learning_rate)))
# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(params_to_update, lr=learning_rate, momentum=0.9)
# Setup the loss fxn
criterion = nn.CrossEntropyLoss()
# Train and evaluate
model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs)
print("start finetuning")
unfreeze(model_ft, args.unfreeze_layer)
# Observe that all parameters are being optimized
finetune_optimizer_ft = optim.SGD(params_to_update, lr=finetune_learning_rate, momentum=0.9)
# Setup the loss fxn
finetune_criterion = nn.CrossEntropyLoss()
model_ft, hist = train_model(model_ft, dataloaders_dict, finetune_criterion, finetune_optimizer_ft, num_epochs=num_epochs, init_acc = hist[-1])
writer.flush()
writer.close()
# Save model
torch.save(model_ft, os.path.join(os.path.join(args.base_path, 'model') + "batch_size" + str(learning_rate) + "_lr_" + str(finetune_learning_rate)))
del model_ft, criterion, optimizer_ft
torch.cuda.empty_cache()I pass in "0" for gpu for the first program, and "1" for the gpu for the second program.
nvidia-smi status after running two separate programs
EDIT : It seems like our CPU usage reaches 99.9% if we run the above code. Would there be any suggestion how to optimize? Thanks!