ResNet50 Training Based on KSpeed - Intelligent Computing LINGJUN

KSpeed accelerates image data loading for computer vision training. This topic walks through integrating KSpeed with a ResNet50 model based on NVIDIA's open-source DeepLearningExamples, using a git patch for the required code modifications. For details, see Key Module Integration of KSpeed .

Code Preparation

Code base repository:

https://github.com/NVIDIA/DeepLearningExamples/commit/174b3d40bfc26f2adcf252676d38d6d5ffa7cbdc

git clone https://github.com/NVIDIA/DeepLearningExamples.git

cd DeepLearningExamples

git checkout master

git reset --hard 174b3d40bfc26f2adcf252676d38d6d5ffa7cbdc

Integrate KSpeed Code

#Stay in the DeepLearningExamples directory
wget http://kspeed-release.oss-cn-beijing.aliyuncs.com/kspeed_resnet50.patch

git apply kspeed_resnet50.patch

Environment configuration

Run the following command to start the training container:

docker run -it --gpus all --name=resnet50_kspeed_test --net=host --ipc host --device=/dev/infiniband/ --ulimit memlock=-1:-1 -v /{path-to-imagenet}:/{path-to-imagenet-in-docker} -v /{path-to-DeepLearningExamples}:/{path-to-DeepLearningExamples-in-docker} eflo-registry.cn-beijing.cr.aliyuncs.com/eflo/ngc-pytorch-kspeed-22.05-py38:v2.2.0

Note

In the above command

{path-to-imagenet} represents the path to the imagenet dataset on the physical machine;
{path-to-imagenet-in-docker} represents the path where the user maps the dataset into the container;
{path-to-DeepLearningExamples} indicates the path to the model training code on the physical machine;
{path-to-DeepLearningExamples-in-docker} indicates the path where the model training code is mapped into the container;

Set these paths according to your environment.

The ImageNet dataset directory structure:

imagenet
├── train
│   ├── n01440764
│   │  ├── n01440764_10026.JPEG
│   │  ├── n01440764_10027.JPEG
│   │  └── ......
│   ├── n01443537
│   └── ......         
└── val                
    ├── n01440764
    │  ├── ILSVRC2012_val_00000293.JPEG
    │  ├── ILSVRC2012_val_00002138.JPEG
    │  └── ......
    ├── n01443537
    └── ......

Dataset acquisition method reference: https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/Classification/ConvNets/resnet50v1.5

Run model training

#Stay in the DeepLearningExamples directory
cd ./PyTorch/Classification/ConvNets

#Single-machine eight-GPU baseline
bash ./resnet50v1.5/training/AMP/DGXA100_resnet50_AMP_multi.sh pytorch {path-to-imagenet-in-docker}

#Single-machine eight-GPU kspeed
bash ./resnet50v1.5/training/AMP/DGXA100_resnet50_AMP_multi.sh kspeed {path-to-imagenet-in-docker}

#Single-machine eight-GPU dali+kspeed
bash ./resnet50v1.5/training/AMP/DGXA100_resnet50_AMP_multi.sh dali-kspeed {path-to-imagenet-in-docker}

Note

In the above command

{path-to-imagenet-in-docker} represents the path of the imagenet dataset in the container, which must be consistent with the path set when starting the container.

Before executing the KSpeed test, ensure that the kspeed service has been deployed.

Key Module Description for KSpeed Integration

Add the kspeeddataloader module file

Add the file DeepLearningExamples/PyTorch/Classification/ConvNets/image_classification/kspeeddataloader.py, which implements a KSpeed-based PyTorch Dataloader and a KSpeed-based Dali Dataloader.

Pytorch Dataloader Based on KSpeed

To implement a KSpeed-based PyTorch Dataloader, modify the Dataset and integrate it with PyTorch's native Sampler and Dataloader:

Import the kspeeddataset module

import kspeed.utils.data.kspeeddataset as KSpeedDataset

Replace torchvison.datasets.ImageFolder with KSpeedDataset.KSpeedImageFolder to enable KSpeed data loading acceleration

train_dataset = KSpeedDataset.KSpeedImageFolder(
        traindir, None, workers, kspeed_iplist,
        "admin", "admin", transforms.Compose(transforms_list),
    )

val_dataset = KSpeedDataset.KSpeedImageFolder(
        valdir, None, workers, kspeed_iplist,
        "admin", "admin",
        transforms.Compose(
            [
                transforms.Resize(
                    image_size + crop_padding, interpolation=interpolation
                ),
                transforms.CenterCrop(image_size),
            ]
        ),
    )

Implement the get_kspeed_train_loader and get_kspeed_val_loader methods. For details, see kspeeddataloader.py lines 16–72 and 74–128

KSpeed-based Dali Dataloader

To implement a KSpeed-based Dali Dataloader, set the Dali pipeline's input data source to KSpeedCallable:

KSpeedCallable

KSpeedCallable inherits KSpeedDataset.KSpeedFolder and reads ImageNet dataset samples through self.dataset.getBIN(path) (see kspeeddataloader.py lines 164–179).

def __call__(self, sample_info):
        
    if self.dataset is None:
        self.load()
    if sample_info.iteration >= self.full_iters:
        raise StopIteration()
    if self.last_seen_epoch != sample_info.epoch_idx:
        self.last_seen_epoch = sample_info.epoch_idx
        self.perm = np.random.default_rng(seed=42 + sample_info.epoch_idx).permutation(len(self.files))
    idx = self.perm[sample_info.idx_in_epoch + self.shard_offset]
        
    path = os.path.join(self.root, self.files[idx])
    dout = self.dataset.getBIN(path)
    sample = np.frombuffer(dout, dtype=np.uint8)
    label = np.int32([self.labels[idx]])
    return sample, label

Dali Pipeline Based on KSpeedCallable

Lines 223–229 of kspeeddataloader.py use KSpeedCallable as the external data source for the Dali Pipeline:

if kspeed:
    images, labels = fn.external_source(source=kscallable,
                        num_outputs=2,
                        batch=False, 
                        parallel=True, 
                        dtype=[types.UINT8, types.INT32], 
                        device='cpu')

Add DATA_BACKEND_CHOICES options

In line 40 of DeepLearningExamples/PyTorch/Classification/ConvNets/image_classification/dataloaders.py, modify the original DATA_BACKEND_CHOICES = ["pytorch", "synthetic"] as follows:

DATA_BACKEND_CHOICES = ["pytorch", "syntetic", "kspeed", "dali-kspeed",  "dali"]

Add args.data_backend option

In DeepLearningExamples/PyTorch/Classification/ConvNets/main.py, add the following code to the args.data_backend branch in lines 512–520:

elif args.data_backend == "kspeed":
    get_train_loader = get_kspeed_train_loader
    get_val_loader = get_kspeed_val_loader
elif args.data_backend == "dali":
    get_train_loader = get_dali_kspeed_train_loader(dali_cpu=True, kspeed=False)
    get_val_loader = get_dali_kspeed_val_loader(dali_cpu=True, kspeed=False)
elif args.data_backend == "dali-kspeed":
    get_train_loader = get_dali_kspeed_train_loader(dali_cpu=True)
    get_val_loader = get_dali_kspeed_val_loader(dali_cpu=True)