Okay, so now I have done some runs with the following code:
import os

from tqdm import tqdm
from torchvision.models import resnet50, resnet18
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import torch
import torch.nn as nn

n = 1000000

class RandomDs(Dataset):
    def __init__(self, ):
        pass

    def __len__(self):
        return n

    def __getitem__(self, index):
        return torch.rand(3, 512, 256)


if __name__ == '__main__':

    dataset = RandomDs()
    data_loader = DataLoader(dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory=False)


    model = resnet50()

    # os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    device = torch.device('cuda:0')
    model = nn.DataParallel(model)
    
    model.to(device)

    with torch.no_grad():
        for batch in tqdm(data_loader):
            batch = batch.to(device)
            model(batch)


I changed the input shape to the network, batch size, number of workers, model type, torch version, and whether it runs with DataParallel or as two scripts (so two scripts, running at the same time, so that both GPUs are in use). These are the results:




input shape
batch_size
num_workers
model
parallel
torch version
crash
memory (GB each)




3, 256, 256
128
4
resnet50
yes
1.2.0
No
~2.3


3, 512, 256
128
4
resnet50
yes
1.2.0
Yes
~3.3


3, 256, 256
256
4
resnet50
yes
1.2.0
Yes
~3.3


3, 256, 256
256
4
resnet50
two scripts
1.2.0
No
~5.3 (single)


3, 256, 256
256
16
resnet50
yes
1.2.0
Yes
~3.3


3, 256, 256
256
4
resnet18
yes
1.2.0
Yes
~2.2


3, 512, 256
128
4
resnet50
yes
1.4.0
Yes
~3.3



It seems to be memory related, but given that I am using 2*2080Ti and I have 64GB ram, it’s not an OOM. Any ideas what else I can test?
I am using cuda 10.2 and tried pytroch version 1.2 and 1.4