import torch
import torchvision.models as models
model = models.resnet18(pretrained=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = torch.nn.DataParallel(model)
model.to(device)
loss_criterion = torch.nn.CrossEntropyLoss()
inp = torch.randn(5, 3, 224, 224).to(device)
target = torch.empty(5, dtype=torch.long).random_(1000).to(device)
output = model(inp)
loss = loss_criterion(output, target)
loss.backward()
I observed that backward fails when I give the input with batch size less than 5. Further if I am using nn.dataparallel, the model also requires there are atleaset 5 samples per batch other wise it gives an error.
I am attaching the screenshot of the error below.
I have tested with both cases, using a single GPUs and multiple GPUs. In both the cases, I have found that the minimum data points in each GPU to be 5 for correctly backpropagating.
Collecting environment information…
PyTorch version: 1.8.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A
OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti
Nvidia driver version: 455.32.00
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.0.5
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7
/usr/local/cuda-9.1/targets/x86_64-linux/lib/libcudnn.so.7
HIP runtime version: N/A
MIOpen runtime version: N/A
