Train arcgis.learn models on multiple GPUs | ArcGIS API for Python

Accelerating AI with GPUs

In this guide, we walk you through how arcgis.learn models with PyTorch backend, support training across multiple GPUs. You can verify how many GPUs are available to PyTorch by running the following commands:

import torch print ('Available devices ', torch.cuda.device_count())

PyTorch provides the capability to utilize multiple GPUs in two ways:

Data Parallelism
Distributed Data Parallelism

arcgis.learn uses one of these two methods to train models using multiple GPUs. Each method has its own significance and both offer an easy way to wrap your code to enable training on multiple GPUs.

Data Parallelism: Data Parallelism refers to splitting the mini-batch of samples into multiple smaller mini-batches and running the computation for each of the smaller mini-batches in parallel across multiple GPUs on a single machine.

Data Parallelism is implemented using torch.nn.DataParallel. You can wrap a module in DataParallel, and it will be parallelized over multiple GPUs in the batch dimension. For more details click here.

For certain models, arcgis.learn already provide support for data parallelism to enhance model performance. This makes it easy for users to utilize multiple GPUs while training a model on a single machine. Below is a list of the models that have DataParallel support.

FeatureClassifier
SingleShotDetector

You can set a subset of GPUs to be used for training your model by running the following command in the first cell:

import os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' # for setting the first 2 GPUs

Distributed Data Parallelism (DDP): In DDP, data is partitioned across multiple devices or nodes (such as GPU-equipped instances). It is implemented using torch.nn.DistributedDataParallel. Unlike data parallelism, DDP supports model parallelism across multiple machines, making it significantly more scalable.

This approach parallelizes by chunking data along the batch dimension and distributing it across specified devices. Ideally, one process is spawned for each model replica; however, a single replica may span multiple devices if the model is large. During the backward pass, DDP automatically synchronizes gradients across all processes to ensure consistent model updates. To read more, click here. Currently, the following models in arcgis.learn support DistributedDataParallel:

UnetClassifier
DeepLab
PSPNetClassifier
MaskRCNN
FasterRCNN
HEDEdgeDetector
BDCNEdgeDetector
ModelExtension
SuperResolution
DETReg
MaXDeepLab
MMDetection
MMSegmentation
RTDetrV2
SamLoRA

The following blocks of code download a sample script that you can use to test the functionality of multi-GPU support, provided you have multiple GPUs. You can modify the script downloaded at script_path to include your model and dataset.

from arcgis.gis import GIS

gis = GIS()

script_item = gis.content.get("afd1c9a88a6f4f04896b4172c0f3a78c")

script_item.name

'train_model.zip'

filepath = script_item.download(file_name=script_item.name)

import zipfile
import os
from pathlib import Path

with zipfile.ZipFile(filepath, "r") as zip_ref:
    zip_ref.extractall(Path(filepath).parent)

script_path = Path(os.path.join(os.path.splitext(filepath)[0]) + ".py")

script_path

WindowsPath('C:/Users/Admin/AppData/Local/Temp/train_model.py')

To run multiple GPUs while training your model, you can download the script from above and/or create one for your own training data. Execute the command shown below in your command prompt:

python -m torch.distributed.launch --nproc_per_node=2 train_model.py

nproc_per_node = number of GPU instances to be used on the given machine.

For detailed arguments of DDP (Distributed Data Parallel), please refer to this page.

To verify that your GPUs are being utilized for training, run nvidia-smi as shown below:

References: