Hello there. I’m new to PyTorch and I’m trying to utilize the ENet NN (from github 1) to classify underwater images from SUIM dataset and evaluate its performance. Through this post I tried to implement the class index mapping of the masks but when I start training the network, some specific masks cause a crash on cuda. I’ve managed to isolate some masks that cause this strange behaviour, as described below.
Follow the main snippets and code adapted from the previously mentioned post:
import torch.nn.functional as F
from collections import OrderedDict
from train import Train
from test import Test
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms.functional as TF
from torch.utils.data import Dataset, DataLoader
import torchvision
import torchvision.transforms as tvtransforms
import transforms as ext_transforms
import torch.optim.lr_scheduler as lr_scheduler
import os
import numpy as np
import matplotlib.pyplot as plt
import glob
from PIL import Image
from enet import ENet
from iou import IoU
import utils
from torchvision import models

class CustomDataset(Dataset):
    def __init__(self, image_paths, target_paths, transform=None, train=True):
        self.image_paths = image_paths
        self.target_paths = target_paths
        self.transform = transform
        self.img = None
        self.msk = None
        self.mapping = {
            (0, 0, 0) : 0, # 'Background'
            (0, 0, 255) : 1, # 'Human Divers'
            (0, 255, 0) : 2, # 'Aquatic Plants and Sea-Grass'
            (0, 255, 255) : 3, # 'Wrecks and Ruins'
            (255, 0, 0): 4, # 'Robots'
            (255, 0, 255) : 5, # 'Reefs and Intertebrates'
            (255, 255, 0) : 6, # 'Fishs and Vertebrates'
            (255, 255, 255) : 7 # 'Sea-Floor and Rocks'
        }
        
    def mask_to_class_rgb(self, mask):
        mask = np.array(mask)
        mask = mask[...,:3]
        mask = torch.from_numpy(mask)
        mask = torch.squeeze(mask)

        print('Unique values in rgb: ', torch.unique(mask))
        class_mask = mask
        class_mask = class_mask.permute(2, 0, 1).contiguous()
        h, w = class_mask.shape[1], class_mask.shape[2]
        mask_out = torch.empty(h, w, dtype=torch.long)
        
        for i in self.mapping:
            idx = (class_mask == torch.tensor(i, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))
            validx = (idx.sum(0) == 3)
            mask_out[validx] = torch.tensor(self.mapping[i], dtype=torch.long)
            
        print('Unique values mapped: ', torch.unique(mask_out))
        
        return mask_out
    

    def __getitem__(self, index):
        img = Image.open(self.image_paths[index]).convert("RGB")
        msk = Image.open(self.target_paths[index])
                
        img_new = tvtransforms.Resize((std_size, std_size))(img)
        mask_new = tvtransforms.Resize((std_size, std_size),  Image.NEAREST)(msk)

        if self.transform is not None:    
            
            img_new = self.transform(img_new)
            
            img_new = img_new.float()
        else:
            img_new = img_new.float()
            
                
        norm = tvtransforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        img_new = norm(img_new)
        
        mask_new = self.mask_to_class_rgb(mask_new)
        
        return img_new, mask_new
        
    
    def __len__(self):
        return len(self.image_paths)

This is the training snippet:
best_miou = 0
if 'start_epoch' in globals():
    start_epoch = start_epoch
else:
    start_epoch = 0

for epoch in range(start_epoch, num_epochs):
        print(">>>> [Epoch: {0:d}] Training".format(epoch))
        
        my_lr_scheduler.step()

        epoch_loss, (iou, miou) = train.run_epoch(True) # Decides whether to print loss at each step

        print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
              format(epoch, epoch_loss, miou))

        if (epoch + 1) % 10 == 0 or epoch + 1 == num_epochs:
            print(">>>> [Epoch: {0:d}] Validation".format(epoch))

            loss, (iou, miou) = val.run_epoch(True)

            print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
                  format(epoch, loss, miou))

            # Print per class IoU on last epoch or if best iou
            if epoch + 1 == num_epochs or miou > best_miou:
                for key, class_iou in zip(class_encoding.keys(), iou):
                    print("{0}: {1:.4f}".format(key, class_iou))

            # Save the model if it's the best thus far
            if miou > best_miou:
                print("\nBest model thus far. Saving...\n")
                best_miou = miou
                utils.save_checkpoint_mod(net, optimizer, epoch + 1, best_miou, 'teste', 'save')

Here is the output:
>>>> [Epoch: 0] Training
Unique values in rgb:  tensor([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
         14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
         28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
         42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  55,  56,
         59,  60,  61,  62,  63,  64,  65,  66,  67,  69,  70,  72,  73,  74,
         75,  76,  80,  81,  83,  84,  86,  87,  91,  93,  94,  95,  97, 119,
        161, 162, 166, 167, 170, 171, 172, 175, 176, 177, 178, 179, 182, 183,
        184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197,
        198, 199, 200, 201, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214,
        215, 216, 217, 218, 219, 221, 222, 223, 224, 225, 226, 227, 228, 229,
        230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243,
        244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255],
       dtype=torch.uint8)
Unique values mapped:  tensor([                0, 42915137138293276, 43204295812155417,
         ..., 66953201512118048, 67517152192928824,
        67519561669712197])
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-59-80f00525c548> in <module>()
     10         my_lr_scheduler.step()
     11 
---> 12         epoch_loss, (iou, miou) = train.run_epoch(True) # Decides whether to print loss at each step
     13 
     14         print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".

4 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2218         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2219     elif dim == 4:
-> 2220         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2221     else:
   2222         # dim == 3 or dim > 4

IndexError: Target 52516180046732819 is out of bounds.

I’m attaching the isolated mask that causes this behaviour. Other masks that are similar to this one cause equal outcome. Also, I’ve edited this mask and made it entirely black, which didn’t crash during training.
d_r_112_1280×720 20.7 KB
Any help would be greatly appreciated!