Время прочтения: 3 мин.

Важным для качественного решения задач CV (Computer Vision) с помощью нейронных сетей, помимо наличия качественной модели (зачастую уже предобученной на других задачах), также является датасет с достаточным количеством изображений (несколько десятков тысяч). Получить необходимый объем размеченных изображений зачастую довольно затруднительно, тогда на помощь может прийти аугментация. Аугментация позволяет увеличить объем исходного количества изображений за счет их изменений: поворот, растягивание/сжатие, изменение цветов и т.д.

Различные методы аугментации ранее уже были разобраны https://newtechaudit.ru/uvelichivaem-dataset-s-pomoshhyu-albumentations-dlya-cv/ другим автором, я же покажу как его можно применить при создании своего даталодера (DataLoader).

Для сокращения времени процесса обучения нейронных сетей используют графические ускорители (GPU), объем памяти которых не способен вместить одновременно весь датасет и обучаемую модель. Для решения этой проблемы используют DataLoader, который «скармливает» нейросети данные из датасета порционно (батчами).

И кажется, что нет проблем: взять готовые архитектуры для аугментации, применить к датасету и поместить в даталодер. Однако, на данный момент DataLoader и Dataset в Pytorch не работают «из коробки» с популярной библиотекой для аугментации albumentations.

Выходом из этого является написание собственного класса Dataset. В данном случае — это Dataset для изображений Imagefolder (структура хранения изображений, при которой каждый класс хранится в папке с соответствующим именем). Для работы понадобится импорт следующих библиотек:

import os
import albumentations as A
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader

Создадим свой собственный класс ImageFolder, наследуя из класса Dataset:
class ImageFolder(Dataset):
    def __init__(self, root_dir, transform=None, total_classes=None):
        self.transform = transform
        self.data = []
        
        if total_classes:
            self.classnames  = os.listdir(root_dir)[:total_classes] # for test
        else:
            self.classnames = os.listdir(root_dir)
            
        for index, label in enumerate(self.classnames):
            root_image_name = os.path.join(root_dir, label)
            
            for i in os.listdir(root_image_name):
                full_path = os.path.join(root_image_name, i)
                self.data.append((full_path, index))
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        data, target = self.data[index]
        img = np.array(Image.open(data))
        
        if self.transform:
            augmentations = self.transform(image=img)
            img = augmentations["image"]
        
        target = torch.from_numpy(np.array(target))
        img = np.transpose(img, (2, 0, 1))
        img = torch.from_numpy(img)            
        
        return img, target

Далее создадим правило, по которому исходное изображение будет меняться:

SIZE = 244
SIZE2 = 256

train_transform_alb = A.Compose(
    [
        A.Resize(SIZE2, SIZE2),
        A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),
        A.RandomCrop(SIZE, SIZE),
        A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),
        A.RandomBrightnessContrast(p=0.5),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ]
)

В данном случае каждое изображение с какой-то долей вероятности (p) поворачивается, сжимается, обрезается, меняет цвета и яркость. А еще все изображения приводятся к одному размеру, а также нормализуются.

Однако, если мы применим трансформацию к исходным данным, их объем не изменится относительно изначальных, поэтому нужно отдельно написать шаги трансформации для исходных данных без аугментации (остаются: приведение к исходному размеру, центрирование и нормализация).

train_transform_base = A.Compose(
    [
        A.Resize(SIZE2, SIZE2),
        A.CenterCrop(SIZE, SIZE),
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ]
)

Далее создадим два датасета и применим к ним подготовленные трансформации:

alb_set = ImageFolder(train_set_dir, transform=train_transform_alb)
base_set = ImageFolder(train_set_dir, transform=train_transform_base)

Объединим датасеты:

full_dataset = alb_set + base_set

Разобьем датасет на обучающий и валидационный:

train_dataset, val_dataset = random_split(full_dataset, [int(0.7*len(full_dataset)), 
                             len(full_dataset) - int(0.7*len(full_dataset))])

Создадим даталодеры для валидационной и обучающего датасетов:

train_loaders = DataLoader(train_dataset, batch_size=16,                                               shuffle=True, num_workers=2)
val_loaders = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=2)

Здесь batch_size – количество картинок в батче, shuffle – перемешивание данных (чтобы лучше происходил процесс обучения, для валидации обычно False), а также num_workers – количество воркеров/ параллельных процессов, позволяет в несколько потоков в процессе обучения модели загружать из диска или оперативной памяти информацию и размещать на графическом ускорителе (что сокращает процесс обучения).

Для num_workers рекомендуется выбирать число не большее, чем количество потоков процессора.

Чтобы итерироваться по батчам, нужно воспользоваться конструкцией:

for xb, yb in train_loaders:
xb, yb = xb.to(device), yb.to(device)

Также можно использовать:

dataiter = iter(check_loaders)
images, labels = dataiter.next()

Далее остается создать свою нейросеть и обучить на данных из даталодера. Обязательно попробуйте!