-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
rocAL - Add pytorch training example (#121)
* Adding training example to rocAL * Update readme * Update readme command to use rocAL dataloader * Update readme for ckpt removal instructions * Resolving review comments * Formatted and moved train script to docs
- Loading branch information
1 parent
e3e6936
commit 366a407
Showing
2 changed files
with
800 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,113 @@ | ||
# ImageNet training in PyTorch | ||
|
||
This example implements training of popular model architectures, such as ResNet, AlexNet, and VGG on the ImageNet dataset. | ||
This version has been modified to use rocAL. It assumes that the dataset is raw JPEGs from the ImageNet dataset. If offers CPU and GPU based pipeline for rocAL - use `rocal-cpu` argument to enable CPU and use `rocal-gpu` argument to enable GPU. | ||
|
||
## Requirements | ||
|
||
- Install PyTorch for [ROCm](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/3rd-party/pytorch-install.html) | ||
- Install rocAL for running rocAL trainings | ||
- Download the ImageNet dataset from http://www.image-net.org/ and use [the following shell script](https://github.com/pytorch/examples/blob/main/imagenet/extract_ILSVRC.sh) to move and extract the training and validation images to labeled subfolders | ||
|
||
## Training | ||
|
||
To train a model, run `imagenet_training.py` with the desired model architecture and the path to the ImageNet dataset: | ||
|
||
```bash | ||
python imagenet_training.py -a resnet18 [imagenet-folder with train and val folders] | ||
``` | ||
|
||
The default learning rate schedule starts at 0.1 and decays by a factor of 10 every 30 epochs. This is appropriate for ResNet and models with batch normalization, but too high for AlexNet and VGG. Use 0.01 as the initial learning rate for AlexNet or VGG: | ||
|
||
```bash | ||
python imagenet_training.py -a alexnet --lr 0.01 [imagenet-folder with train and val folders] | ||
``` | ||
To run a rocAL integrated training, use `rocal-cpu` or `rocal-gpu` | ||
|
||
```bash | ||
python3 imagenet_training.py -a resnet50 -j$(nproc) --batch-size 1024 --rocal-cpu [imagenet-folder with train and val folders] | ||
``` | ||
|
||
Make sure to remove older checkpoints (`rm *.pth.tar`) saved in the folder if the example has been run before | ||
|
||
## Use Dummy Data | ||
|
||
ImageNet dataset is large and time-consuming to download. To get started quickly, run `imagenet_training.py` using dummy data by "--dummy". It's also useful for training speed benchmark. Note that the loss or accuracy is useless in this case. | ||
|
||
```bash | ||
python imagenet_training.py -a resnet18 --dummy | ||
``` | ||
|
||
## Multi-processing Distributed Data Parallel Training | ||
|
||
You should always use the NCCL backend for multi-processing distributed training since it currently provides the best distributed training performance. | ||
|
||
### Single node, multiple GPUs: | ||
|
||
```bash | ||
python imagenet_training.py -a resnet50 --dist-url 'tcp://127.0.0.1:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 1 --rank 0 [imagenet-folder with train and val folders] | ||
``` | ||
|
||
### Multiple nodes: | ||
|
||
Node 0: | ||
|
||
```bash | ||
python imagenet_training.py -a resnet50 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 2 --rank 0 [imagenet-folder with train and val folders] | ||
``` | ||
|
||
Node 1: | ||
|
||
```bash | ||
python imagenet_training.py -a resnet50 --dist-url 'tcp://IP_OF_NODE0:FREEPORT' --dist-backend 'nccl' --multiprocessing-distributed --world-size 2 --rank 1 [imagenet-folder with train and val folders] | ||
``` | ||
|
||
## Usage | ||
|
||
```bash | ||
usage: imagenet_training.py [-h] [-a ARCH] [-j N] [--epochs N] [--start-epoch N] [-b N] [--lr LR] [--momentum M] [--wd W] [-p N] [--resume PATH] [-e] [--pretrained] [--world-size WORLD_SIZE] [--rank RANK] | ||
[--dist-url DIST_URL] [--dist-backend DIST_BACKEND] [--seed SEED] [--gpu GPU] [--multiprocessing-distributed] [--dummy] | ||
[DIR] | ||
|
||
PyTorch ImageNet Training | ||
|
||
positional arguments: | ||
DIR path to dataset (default: imagenet) | ||
|
||
optional arguments: | ||
-h, --help show this help message and exit | ||
-a ARCH, --arch ARCH model architecture: alexnet | convnext_base | convnext_large | convnext_small | convnext_tiny | densenet121 | densenet161 | densenet169 | densenet201 | efficientnet_b0 | | ||
efficientnet_b1 | efficientnet_b2 | efficientnet_b3 | efficientnet_b4 | efficientnet_b5 | efficientnet_b6 | efficientnet_b7 | googlenet | inception_v3 | mnasnet0_5 | mnasnet0_75 | | ||
mnasnet1_0 | mnasnet1_3 | mobilenet_v2 | mobilenet_v3_large | mobilenet_v3_small | regnet_x_16gf | regnet_x_1_6gf | regnet_x_32gf | regnet_x_3_2gf | regnet_x_400mf | regnet_x_800mf | | ||
regnet_x_8gf | regnet_y_128gf | regnet_y_16gf | regnet_y_1_6gf | regnet_y_32gf | regnet_y_3_2gf | regnet_y_400mf | regnet_y_800mf | regnet_y_8gf | resnet101 | resnet152 | resnet18 | | ||
resnet34 | resnet50 | resnext101_32x8d | resnext50_32x4d | shufflenet_v2_x0_5 | shufflenet_v2_x1_0 | shufflenet_v2_x1_5 | shufflenet_v2_x2_0 | squeezenet1_0 | squeezenet1_1 | vgg11 | | ||
vgg11_bn | vgg13 | vgg13_bn | vgg16 | vgg16_bn | vgg19 | vgg19_bn | vit_b_16 | vit_b_32 | vit_l_16 | vit_l_32 | wide_resnet101_2 | wide_resnet50_2 (default: resnet18) | ||
-j N, --workers N number of data loading workers (default: 4) | ||
--rocal-cpu use rocAL CPU dataloader | ||
--rocal-gpu use rocAL GPU dataloader | ||
--epochs N number of total epochs to run | ||
--start-epoch N manual epoch number (useful on restarts) | ||
-b N, --batch-size N mini-batch size (default: 256), this is the total batch size of all GPUs on the current node when using Data Parallel or Distributed Data Parallel | ||
--lr LR, --learning-rate LR | ||
initial learning rate | ||
--momentum M momentum | ||
--wd W, --weight-decay W | ||
weight decay (default: 1e-4) | ||
-p N, --print-freq N print frequency (default: 10) | ||
--resume PATH path to latest checkpoint (default: none) | ||
-e, --evaluate evaluate model on validation set | ||
--pretrained use pre-trained model | ||
--world-size WORLD_SIZE | ||
number of nodes for distributed training | ||
--rank RANK node rank for distributed training | ||
--dist-url DIST_URL url used to set up distributed training | ||
--dist-backend DIST_BACKEND | ||
distributed backend | ||
--seed SEED seed for initializing training. | ||
--gpu GPU GPU id to use. | ||
--multiprocessing-distributed | ||
Use multi-processing distributed training to launch N processes per node, which has N GPUs. This is the fastest way to use PyTorch for either single node or multi node data parallel | ||
training | ||
--dummy use fake data to benchmark | ||
|
||
``` |
Oops, something went wrong.