Why my gpu-util is low?

[yllgl]

I use one node and 4 gpus for training. And I use dali dataloader, I don't know why my gpu util is low, and training is also slow. About 1:30 per epoch, I train for 200 epoches, which will cost 5 hours. It's slower than the project mmclassification, which only cost 3.5 hours. Compared to mmclassification project which can only support torch.utils.data.dataloader, I think if I use dali_dataloader, it will accelerate my training. But as you can see, it's the opposite. I don't know why. Could anyone give me some advice? I use cifar10 dataset. And I train on slurm.

Here is my code.

main.py

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from net import ResNet18
if __name__ == '__main__':
    model = ResNet18()
    trainer = pl.Trainer( max_epochs=200,log_every_n_steps=1,
        log_gpu_memory='min_max',gpus=4,num_nodes=1,accelerator='ddp',
        fast_dev_run=False,callbacks=[ModelCheckpoint(monitor='val_accuracy',mode='max')],
        progress_bar_refresh_rate=1,replace_sampler_ddp=False)
    trainer.fit(model)

net.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from dataloader import dali_DataLoader,HybridPipe,dali_CIFAR10
class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(
            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion *
                               planes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(self.expansion*planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(pl.LightningModule):
    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 64

        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.linear = nn.Linear(512*block.expansion, num_classes)
        self.correct = 0
        self.total_size = 0
    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out
    def training_step(self, batch, batch_idx):
        x, y = batch
        x = self(x)
        loss_fn = nn.CrossEntropyLoss()
        loss = loss_fn(x,y)
        predicted = torch.argmax(x, dim=1, keepdim=False)
        self.correct += (predicted == y).sum().item()
        self.total_size += y.size(0)
        self.log('train_loss', loss,prog_bar=True, logger=True)
        self.log('train_accuracy', self.correct/self.total_size,prog_bar=True, logger=True)
        return loss
    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = self(x)
        loss_fn = nn.CrossEntropyLoss()
        loss = loss_fn(x,y)
        predicted = torch.argmax(x, dim=1, keepdim=False)
        self.correct += (predicted == y).sum().item()
        self.total_size += y.size(0)
        self.log('val_loss', loss,on_step=False, on_epoch=True,prog_bar=True, logger=True)
        self.log('val_accuracy', self.correct/self.total_size,prog_bar=True, logger=True)
        return loss
    def validation_epoch_end(self,out):
        self.log('val_accuracy', self.correct/self.total_size,prog_bar=True, logger=True)
        self.correct=0
        self.total_size=0
    def train_epoch_end(self,out):
        self.log('train_accuracy', self.correct/self.total_size,prog_bar=True, logger=True)
        self.correct=0
        self.total_size=0
    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), lr=0.1)
        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [100,150], gamma=0.1, last_epoch=-1, verbose=False)
        return [optimizer],[scheduler]
    def train_dataloader(self):
        loader = dali_DataLoader(pipelines=HybridPipe(dali_CIFAR10(root='./data'), batch_size=32, pad_ratio=1.25,num_threads=4,
             is_distribute=True, crop_size=32,ramdom_flip=True,
             normalize=dict(mean=[125.307, 122.961, 113.8575],std=[51.5865, 50.847, 51.255])))
        return loader

    def val_dataloader(self):
        loader = dali_DataLoader(pipelines=HybridPipe(dali_CIFAR10(root='./data',test_mode=True), batch_size=100,
             normalize=dict(mean=[125.307, 122.961, 113.8575],std=[51.5865, 50.847, 51.255])))
        return loader
def ResNet18():
    return ResNet(BasicBlock, [2, 2, 2, 2])

dataloader.py

import os,sys,math,random,pickle
import torch
import numpy as np
import torch.distributed as dist
try:
    from nvidia import dali
    from nvidia.dali.pipeline import Pipeline
    import nvidia.dali.types as types
    import nvidia.dali.fn as fn
    import nvidia.dali.ops as ops
    from nvidia.dali.plugin.pytorch import DALIClassificationIterator
except:
    print('Could not import DALI')
class dali_DataLoader():
    def __init__(self, pipelines, **kwargs):
        pipelines.build()
        try:
            self._dali_iterator = DALIClassificationIterator(pipelines=pipelines, size=len(pipelines.iterator.indices))
            self.sampler = pipelines.iterator
        except:
            self._dali_iterator = DALIClassificationIterator(pipelines=pipelines, reader_name='Reader')
            self.sampler = self
    def set_epoch(self,epoch):
        pass
    def __iter__(self):
        return self

    def __len__(self):
        return int(math.ceil(self._dali_iterator._size / self._dali_iterator.batch_size))
    def __next__(self):
        
        try:
            data = next(self._dali_iterator)
        except StopIteration:
            self._dali_iterator.reset()
            raise StopIteration
        # Decode the data output
        input = data[0]['data']
        target = data[0]['label'].squeeze().long()

        return input,target
class identity():
    def __call__(self,x,*tmp,**kargs):
        return x
class HybridPipe(Pipeline):
    def __init__(self,dataset, batch_size, file_root=None,filelist_path=None,num_threads=1, pad_ratio=1,is_distribute=True, resize=None,crop_size=[0,0],ramdom_flip=False,normalize=None,random_rotate_degree=None):
        device_id = torch.cuda.current_device()
        print("device_id",device_id)
        super(HybridPipe, self).__init__(batch_size, num_threads, device_id, seed=12 + device_id)
        
        if is_distribute:
            if filelist_path is not None:
                if file_root is None:
                    raise Exception("if provide filelist_path, then must provide file_root")
                else:
                    self.input = ops.readers.File(file_root=file_root,file_list=filelist_path,num_shards=dist.get_world_size(),prefetch_queue_depth=num_threads,read_ahead=True,shard_id=dist.get_rank())
                    self.decode = ops.decoders.Image(device="mixed", output_type=types.RGB)
                    self.use_file=True
            else:
                self.iterator = iter(Distribute_Input_Iter(dataset, batch_size))
                #self.input = ops.ExternalSource(source=self.iterator, num_outputs=2)
                self.input = ops.ExternalSource()
                self.input_label = ops.ExternalSource()
                self.use_file=False
        else:
            if filelist_path is not None:
                if file_root is None:
                    raise Exception("if provide filelist_path, then must provide file_root")
                else:
                    self.input = ops.readers.File(file_root=file_root,file_list=filelist_path,num_shards=dist.get_world_size(),prefetch_queue_depth=num_threads,read_ahead=True,shard_id=dist.get_rank())
                    self.decode = ops.decoders.Image(device="mixed", output_type=types.RGB)
                    self.use_file=True
            else:
                self.iterator = iter(Normal_Input_Iter(dataset, batch_size))
                self.input = ops.ExternalSource()
                self.input_label = ops.ExternalSource()
                self.use_file=False
        dali_device = "gpu"
        
        if isinstance(resize,(tuple,list)) and len(resize)==2:
            self.resize = ops.Resize(size=tuple(resize))
        elif isinstance(resize,(int, float)):
            self.resize = ops.Resize(size=tuple(resize,resize))
        else:
            self.resize = identity()
        if normalize is not None and isinstance(normalize,dict):
            self.mean = normalize.get('mean',0)
            self.std = normalize.get('std',1)
        else:
            self.mean = 0
            self.std = 1
        if isinstance(crop_size, (int, float)):
            crop_size = [crop_size,crop_size]
        if (len(crop_size)==2 and (crop_size[0]==0 or crop_size[1]==0)):
            self.crop = identity()
        else:
            self.crop = ops.Crop(device=dali_device, crop_h=crop_size[0], crop_w=crop_size[1])
        if pad_ratio>1:
            self.pad = ops.Paste(device=dali_device, ratio=pad_ratio, fill_value=0)
        else:
            self.pad = identity()
        self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                            dtype=types.FLOAT,
                                            output_layout=types.NCHW,
                                            mean=self.mean,
                                            std=self.std
                                            )
        if ramdom_flip:
            self.coin = ops.random.CoinFlip(probability=0.5)
        else:
            self.coin = lambda :0
        if random_rotate_degree is not None:
            try:
                tmp = math.abs(int(random_rotate_degree))
                self.degree = ops.random.Uniform(range=(-tmp, tmp))
                self.rotate = ops.Rotate()
            except:
                self.degree = lambda :0
                self.rotate = identity()
        else:
            self.degree = lambda :0
            self.rotate = identity()
        
    def iter_setup(self):
        if not self.use_file:
            (images, labels) = self.iterator.__next__()
            self.feed_input(self.jpegs, images, layout="HWC")
            self.feed_input(self.labels, labels)

    def define_graph(self):
        rng = self.coin()
        print()
        if self.use_file:
            self.jpegs,self.labels = self.input(name="Reader")
            self.jpegs = self.decode(self.jpegs)
        else:
            self.jpegs= self.input()
            self.labels = self.input_label()
        output = self.jpegs
        output = self.resize(output)
        output = self.rotate(output, angle=self.degree())
        output = self.pad(output.gpu())
        output = self.crop(output)
        output = self.cmnp(output, mirror=rng)
        return [output, self.labels]
class Distribute_Input_Iter():
    def __init__(self,dataset, batch_size, num_replicas=None,rank=None,shuffle=True,seed=0,drop_last=False):
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
            #num_replicas = 1
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
            #rank = 0
        if rank >= num_replicas or rank < 0:
            raise ValueError(
                "Invalid rank {}, rank should be in the interval"
                " [0, {}]".format(rank, num_replicas - 1))
        self.dataset = dataset
        self.batch_size = batch_size
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        
        # If the dataset length is evenly divisible by # of replicas, then there
        # is no need to drop any data, since the dataset will be split equally.
        if self.drop_last and len(self.dataset) % self.num_replicas != 0:  # type: ignore
            # Split to nearest available length that is evenly divisible.
            # This is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil(
                # `type:ignore` is required because Dataset cannot provide a default __len__
                # see NOTE in pytorch/torch/utils/data/sampler.py
                (len(self.dataset) - self.num_replicas) / self.num_replicas  # type: ignore
            )
        else:
            self.num_samples = math.ceil(len(self.dataset) / self.num_replicas)  # type: ignore
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed
        self.epoch=0
        indices = list(range(len(self.dataset)))  # type: ignore

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            padding_size = self.total_size - len(indices)
            if padding_size <= len(indices):
                indices += indices[:padding_size]
            else:
                indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[:self.total_size]
        assert len(indices) == self.total_size,'len(indices) != self.total_size'

        # subsample
        indices = indices[self.rank:self.total_size:self.num_replicas]
        assert len(indices) == self.num_samples,'len(indices) != self.num_samples'
        self.indices = indices
    def set_epoch(self,epoch):
        self.epoch = epoch
    def __iter__(self):
        self.i = 0
        self.n = len(self.indices)
        return self
    def __next__(self):
        batch = []
        labels = []
        should_shuffle = False
        
        for _ in range(self.batch_size):
            if self.i % self.n == self.n-1:
                should_shuffle = True
            img, label = self.dataset.__getitem__(self.indices[self.i])
            batch.append(img)
            labels.append(label)
            self.i = (self.i + 1) % self.n
        if should_shuffle:
            random.shuffle(self.indices)
        return (batch, labels)
class Normal_Input_Iter():
    def __init__(self,dataset, batch_size):
        self.dataset = dataset
        self.batch_size = batch_size
        self.indices = list(range(len(self.dataset)))
    def __iter__(self):
        self.i = 0
        self.n = len(self.dataset)
        return self
    def __next__(self):
        batch = []
        labels = []
        should_shuffle = False
        
        for _ in range(self.batch_size):
            if self.i % self.n == self.n-1:
                should_shuffle = True
            img, label = self.dataset.__getitem__(self.indices[self.i])
            batch.append(img)
            labels.append(label)
            self.i = (self.i + 1) % self.n
        if should_shuffle:
            random.shuffle(self.indices)
        return (batch, labels)