use gpu in customized pipeline

[ChenJunzhi-buaa]

I try to customize a pipeline. In it，image is cropped by a network，which uses gpu.

However, this pipeline raise errors, about using gpu in multiprocessing, as below:

I tried some method but failed.

What can I do?

[Ezra-Yu]

Please provide more information on how to use pipeline in the GPU and what your config looks like

[ChenJunzhi-buaa]

customized pipeline

Thank you for your help！
I use gpu as below，which is a exemple. Processing images by resnet18 network with gpu.

from ..builder import PIPELINES
import cv2
import torch
import numpy as np
from unet_from_scratch.model.unet import Model
import multiprocessing
import mmcv
import torchvision

@PIPELINES.register_module()
class CropBySeg(object):
    def __init__(self, margin=20, resize_shape=(2048, 2048), interpolation=True,
            model_path='unet_from_scratch/trains_log/OD_20221206-0815_c-32_interpolation-True_max-lr-4e-05_epoch-100/latest.pkl'):
        self.margin = margin
        self.resize_shape = resize_shape
        self.interpolation = interpolation
        self.model = Model(c=32, interpolation=interpolation, local_rank=-1)
        self.model.load_model(path=model_path, rank=-1)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # self.device = 'cpu'
        self.model.unet.to(self.device)

    def __call__(self, results):
        for key in results.get('img_fields', ['img']):
            img = results[key]
            crop_img = self.crop_by_seg(self.model, img, self.margin, self.resize_shape)
            results[key] = crop_img
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'margin{self.margin}' + f'interpolation{self.interpolation}'

    def crop_by_seg(self, model, img, margin=20, resize_shape=(2048, 2048)):
        if resize_shape != img.shape:
            img = cv2.resize(img, resize_shape)  
        
        ##### use gpu like below ######
        input = torch.randn((2,3,32,32)).to(self.device)
        resnet = torchvision.models.resnet18().to(self.device)
        out = resnet(input).to('cpu')

        ##### the complex but not necessary codes are deleted for clarity

        crop_img = img
        return crop_img

my config

# add CropBySeg in train_pipeline and test pipe line
train_pipeline.append(dict(type = 'CropBySeg'))
train=dict(pipeline = train_pipeline)
    

[Ezra-Yu]

with #1270, probably because there are multiple processes involved, the dataloader will open many processes and many UNet will be put into cuda at the same time to calculate, you can try this to reduce the batch size and workers_per_gpu or to place this transform is in the model.data_preprocessor.

[ChenJunzhi-buaa]

I set batch size to be 8, workers_per_gpu to be 1. And training is successful, but after that, val is still OOM.
Maybe I have to try to place this transform is in the model.data_preprocessor.

model.data_preprocesso is not in multiple processes?

Where is model.data_preprocessor？I didn't find.

[Ezra-Yu]

https://github.com/open-mmlab/mmclassification/blob/dev-1.x/mmcls/models/utils/data_preprocessor.py

model.data_preprocessor do :

1. normalization and collate all the images together
2. batch_augments
3. stack the targets together
4. sometimes do target into one-hot

Maybe you can do that at the beginning, you need to process images one by one.

Or at the end, you need to process images after seg and then resize, collate.