High GPU Memory Demand for pytorch?
Category: vision

gaoking132: Hi, I have a variant of inception v3 network, which is pre-trained and fixed. Now I want to do the inference for that network. I tried both tensorflow and pytorch. I don’t know why this happens. My guess is that tensorflow may not cache the intermediate feature maps in the graphdef mode, but pytorch may do. I also had other problems related to the memory usage. My recent experience suggests that pytorch often uses higher gpu memory than tensorflow. In many cases, I have to reduce batch size, which may or may not solve my problem. Are there any suggestions on how to use pytorch more memory efficiently?

jekbradbury: 

apaszke: 

gaoking132: Thanks for the pointer. I think that my problem might be a little different. I have two networks A and B, the output of A will be fed to network B, and the loss is defined on the top of network B. The network B is fixed and network A is to be optimized. So I need network B to back propagate the gradients from loss to A, but these gradients do not need to be saved since network B doesn’t need to be updated. I am wondering that in this case, can I set all variables in B to be volatile=True? Is there some way for me to know whether the gradients of parameters in B are buffered or not?

smth: <SCODE># let's assume input is a Tensor, A and B are networks

optimizerA = optim.SGD(A.parameters(), ...)

# freeze B's params
for p in B.parameters():
    p.requires_grad = False

iv = Variable(input)
optimizerA.zero_grad()
out1 = A(iv)
out2 = B(out1)
out2.backward()
optimizer.step()
<ECODE>

gaoking132: Thanks. This is what I did. One quick question is that if I set requires_grad = False, does it save extra gpu memory for me?

smth: Yes it saves memory at certain places

gaoking132: OK, I felt a little frustrated. I did everything I could to reduce GPU memory. But the usage of my pytorch code is still more than twice memory-consuming than my tensorflow implementation. Because the code is confidential, I cannot release it in public. My network B is an inception v3 network. I found that after I passed the output of A to B, the GPU memory increases 4 GB. The batch size is only about 20.

smth: 

smth: if you can give a small script that showcases your code (without giving out your actual code), happy to take a look.

yunjey: 

apaszke: 

jhjungCode: 

gaoking132: Thank you for all the helpful replies. The code for both pytorch and tensorflow implementation is a little complex, which certainly cannot be used as a fair comparison. I don’t have too much time to extract a small part of the code for debugging purpose. Sorry for not providing helpful feedbacks on this thread.

jhjungCode: This is my test codes for comparing pytorch and tensorflow <SCODE>import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import numpy as np

is_cuda = torch.cuda.is_available() # if cuda is avaible, True
traindir = './flower_photos'

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

batch_size = 16 
train_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(traindir,
                         transforms.Compose([
                             transforms.RandomSizedCrop(224),
                             transforms.RandomHorizontalFlip(),
                             transforms.ToTensor(),
                             normalize,])),
    batch_size=batch_size,
    shuffle=True,
    num_workers=4)

cls_num = len(datasets.folder.find_classes(traindir)[0])

test_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(traindir,
                         transforms.Compose([
                             transforms.RandomSizedCrop(224),
                             transforms.RandomHorizontalFlip(),
                             transforms.ToTensor(),
                             normalize,])),
    batch_size=batch_size,
    shuffle=True,
    num_workers=1)

model = torchvision.models.resnet152(pretrained = True)

### don't update model parameters
for param in model.parameters() :
    param.requires_grad = False
#modify last fully connected layter
model.fc = nn.Linear(model.fc.in_features, cls_num)

fc_parameters = [
    {'params': model.fc.parameters()},
]
optimizer = torch.optim.Adam(fc_parameters, lr=1e-4, weight_decay=1e-4)
loss_fn = nn.CrossEntropyLoss() 

if is_cuda : model.cuda(), loss_fn.cuda()
# trainning

model.train()
train_loss = []
train_accu = []
i = 0
for epoch in range(1):
    for image, target in train_loader:
        image, target = Variable(image.float()), Variable(target) 
        if is_cuda :  image, target = image.cuda(), target.cuda() 
        output = model(image) 
        loss = loss_fn(output, target) 
        optimizer.zero_grad() 
        loss.backward() 
        optimizer.step() 
        
        pred = output.data.max(1)[1]
        accuracy = pred.eq(target.data).sum()/batch_size
        
        train_loss.append(loss.data[0])
        train_accu.append(accuracy)

        if i % 300 == 0:
            print(i, loss.data[0])
        i += 1
<ECODE> <SCODE>images, _, labels = load_batch(dataset,  batch_size=256, height=image_size, width=image_size)
<ECODE>

apaszke: Thanks! We’ll look into that!

apaszke: Thanks for posting the repro!

jhjungCode: 

gaoking132: Thanks, it also solves my problem.

Mario_Parreno: Hello! It is supposed to work as -> torch.no_grad() ? Because with torch.no_grad(): my model inferences consumes the same memory as in train… It is normal?