Repository: pytorch
Issue: NaN Loss for FasterRCNN on Multiclass Object Detection on Custom Dataset COCO
Labels: triaged,module: vision

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce
Code
# Dataloader / DataLoading
```
class OwnDataset(torch.utils.data.Dataset):
    def __init__(self, root, annotation, transforms=None):
        self.root = root
        self.transforms = transforms
        self.coco = COCO(annotation)
        self.ids = list(sorted(self.coco.imgs.keys()))

    def __getitem__(self, index):
        # Own coco file
        coco = self.coco
        # Image ID
        img_id = self.ids[index]
        # List: get annotation id from coco
        ann_ids = coco.getAnnIds(imgIds=img_id)
        # Dictionary: target coco_annotation file for an image
        coco_annotation = coco.loadAnns(ann_ids)
        # print(coco_annotation)
        # path for input image
        path = coco.loadImgs(img_id)[0]['file_name']
        # open the input image
        img = Image.open(os.path.join(self.root, path)).convert("RGB")
        width, height = img.size
        # number of objects in the image
        num_objs = len(coco_annotation)

        # Bounding boxes for objects
        # In coco format, bbox = [xmin, ymin, width, height]
        # In pytorch, the input should be [xmin, ymin, xmax, ymax]
        boxes = []
        labelings = []
        for i in range(num_objs):
            xmin = max(coco_annotation[i]['bbox'][0], 0)
            ymin = max(coco_annotation[i]['bbox'][1], 0)
            xmax = max(xmin + coco_annotation[i]['bbox'][2], 0)
            ymax = max(ymin + coco_annotation[i]['bbox'][3], 0)
            boxes.append([xmin, ymin, xmax, ymax])
            labelings.append(coco_annotation[i]['category_id'])
            assert xmin >= 0
            assert xmax <= width
            assert xmin <= xmax

            assert ymin >= 0
            assert ymax <= height
            assert ymin <= ymax
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        boxes = boxes.reshape(-1, 4)
        # Labels (In my case, I only one class: target class or background)
        labels = torch.as_tensor(labelings, dtype=torch.int64)
        # Tensorise img_id
        img_id = torch.tensor([img_id])
        # Size of bbox (Rectangular)
        areas = []
        for i in range(num_objs):
            areas.append(coco_annotation[i]['area'])
        areas = torch.as_tensor(areas, dtype=torch.float32)
        # Iscrowd
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # Annotation is in dictionary format
        my_annotation = {}
        my_annotation["boxes"] = boxes
        my_annotation["labels"] = labels
        my_annotation["image_id"] = img_id
        my_annotation["area"] = areas
        my_annotation["iscrowd"] = iscrowd

        if self.transforms is not None:
            img = self.transforms(img)
        
        # img.type(torch.float32)

        return img, my_annotation

    def __len__(self):
        return len(self.ids)

from torchvision import transforms
def get_transform():
    custom_transforms = [
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]
    return torchvision.transforms.Compose(custom_transforms)

# path to your own data and coco file
train_data_dir = '/root/output2/imgs/'
train_coco = '/root/output2/train.json'

# create own Dataset
my_dataset = OwnDataset(root=train_data_dir, 
                        annotation=train_coco, 
                        transforms=get_transform()
                        )

# collate_fn needs for batch
def collate_fn(batch):
    return tuple(zip(*batch))

# Batch size
train_batch_size = 2

# own DataLoader
data_loader = torch.utils.data.DataLoader(my_dataset,
                                          batch_size=train_batch_size,
                                          shuffle=True,
                                          num_workers=2,
                                          collate_fn=collate_fn,
                                          pin_memory=True)
```

# Training procedure and Model Instace for FasterRCNN
```
def get_model_instance_segmentation(num_classes):
    backbone = torchvision.models.mobilenet_v2(pretrained=True, progress=False).features

    backbone.out_channels = 1280
    anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256),),
                                    aspect_ratios=((0.5, 1.0, 2.0),))
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
                                                  output_size=7,
                                                  sampling_ratio=2)
    model = FasterRCNN(backbone,
                   num_classes=num_classes,
                   rpn_anchor_generator=anchor_generator,
                   box_roi_pool=roi_pooler)
    return model

model = get_model_instance_segmentation(num_classes)

params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.0001, momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                   step_size=3,
                                                   gamma=0.1)

len_dataloader = len(data_loader)
total_len = num_epochs * len_dataloader
model.to(device)
old_epoch = 0

import time
num_epochs = 10000

time_s = time.time()
total_processed = 0
total_len = num_epochs * len_dataloader
for epoch in range(old_epoch, num_epochs):
    print("---------------------------------------------------------------")
    print("EPOCH: ", epoch)
    model.train()
    for imgs, annotations in data_loader:
        imgs = list(img.to(device) for img in imgs)
        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
        loss_dict = model(imgs, annotations)
        losses = sum(loss for loss in loss_dict.values())
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        lr_scheduler.step()
```


## Expected behavior

I am getting NaN as Loss after some few epochs

## Environment

CUDA used to build PyTorch: 10.2

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 440.64.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.3

Versions of relevant libraries:
[pip3] numpy==1.18.4
[pip3] sagemaker-pytorch-training==1.3.0
[pip3] torch==1.5.0
[pip3] torchvision==0.6.0
[conda] magma-cuda101             2.5.1                         1    pytorch
[conda] mkl                       2019.4                      243    anaconda
[conda] mkl-include               2019.4                      243    anaconda
[conda] mkl-service               2.3.0            py36he904b0f_0    anaconda
[conda] mkl_random                1.1.0            py36hd6b4f25_0    anaconda
[conda] numpy                     1.18.4                   pypi_0    pypi
[conda] sagemaker-pytorch-training 1.3.0                    pypi_0    pypi
[conda] torch                     1.5.0                    pypi_0    pypi
[conda] torchvision               0.6.0                    pypi_0    pypi

## Additional context

<!-- Add any other context about the problem here. -->


cc @fmassa