mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
250 lines
7.4 KiB
Python
250 lines
7.4 KiB
Python
# Copyright 2022 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the \"License\");
|
|
# you may not use this file except in compliance with the License.\n",
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an \"AS IS\" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Train resnet on multiple GPUs with DDP."""
|
|
|
|
import argparse
|
|
import functools
|
|
import itertools
|
|
import math
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
from torch import nn
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
import torchmetrics
|
|
from torchvision.models import resnet50
|
|
from torchvision.transforms import transforms
|
|
import webdataset as wds
|
|
|
|
|
|
def wds_split(src, rank, world_size):
|
|
"""Shards split function for webdataset."""
|
|
# The context of caller of this function is within multiple processes
|
|
# (by DDP world_size) and multiple workers (by dataloader_num_workers).
|
|
# So we totally have (world_size * num_workers) workers for processing data.
|
|
# NOTE: Raw data should be sharded to enough shards to make sure one process
|
|
# can handle at least one shard, otherwise the process may hang.
|
|
worker_id = 0
|
|
num_workers = 1
|
|
worker_info = torch.utils.data.get_worker_info()
|
|
if worker_info:
|
|
worker_id = worker_info.id
|
|
num_workers = worker_info.num_workers
|
|
for s in itertools.islice(src, rank * num_workers + worker_id, None,
|
|
world_size * num_workers):
|
|
yield s
|
|
|
|
|
|
def identity(x):
|
|
return x
|
|
|
|
|
|
def create_wds_dataloader(rank, args, mode):
|
|
"""Create webdataset dataset and dataloader."""
|
|
if mode == 'train':
|
|
transform = transforms.Compose([
|
|
transforms.RandomResizedCrop(224),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
data_path = args.train_data_path
|
|
data_size = args.train_data_size
|
|
batch_size_local = args.train_batch_size
|
|
batch_size_global = args.train_batch_size * args.gpus
|
|
# Since webdataset disallows partial batch, we pad the last batch for train.
|
|
batches = int(math.ceil(data_size / batch_size_global))
|
|
else:
|
|
transform = transforms.Compose([
|
|
transforms.Resize(256),
|
|
transforms.CenterCrop(224),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
data_path = args.eval_data_path
|
|
data_size = args.eval_data_size
|
|
batch_size_local = args.eval_batch_size
|
|
batch_size_global = args.eval_batch_size * args.gpus
|
|
# Since webdataset disallows partial batch, we drop the last batch for eval.
|
|
batches = int(data_size / batch_size_global)
|
|
|
|
dataset = wds.DataPipeline(
|
|
wds.SimpleShardList(data_path),
|
|
functools.partial(wds_split, rank=rank, world_size=args.gpus),
|
|
wds.tarfile_to_samples(),
|
|
wds.decode('pil'),
|
|
wds.to_tuple('jpg;png;jpeg cls'),
|
|
wds.map_tuple(transform, identity),
|
|
wds.batched(batch_size_local, partial=False),
|
|
)
|
|
num_workers = args.dataloader_num_workers
|
|
dataloader = wds.WebLoader(
|
|
dataset=dataset,
|
|
batch_size=None,
|
|
shuffle=False,
|
|
num_workers=num_workers,
|
|
persistent_workers=True if num_workers > 0 else False,
|
|
pin_memory=True).repeat(nbatches=batches)
|
|
print(f'{mode} dataloader | samples: {data_size}, '
|
|
f'num_workers: {num_workers}, '
|
|
f'local batch size: {batch_size_local}, '
|
|
f'global batch size: {batch_size_global}, '
|
|
f'batches: {batches}')
|
|
return dataloader
|
|
|
|
|
|
def train(model, device, dataloader, optimizer):
|
|
model.train()
|
|
for image, target in dataloader:
|
|
image = image.to(device, non_blocking=True)
|
|
target = target.to(device, non_blocking=True)
|
|
pred = model(image)
|
|
# pred.shape (N, C), target.shape (N)
|
|
loss = nn.functional.cross_entropy(pred, target)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
return loss
|
|
|
|
|
|
def evaluate(model, device, dataloader, metric):
|
|
model.eval()
|
|
with torch.no_grad():
|
|
for image, target in dataloader:
|
|
image = image.to(device, non_blocking=True)
|
|
target = target.to(device, non_blocking=True)
|
|
pred = model(image)
|
|
metric.update(pred, target)
|
|
accuracy = metric.compute()
|
|
metric.reset()
|
|
return accuracy
|
|
|
|
|
|
def worker(gpu, args):
|
|
"""Run training and evaluation."""
|
|
# Init process group.
|
|
print(f'Initiating process {gpu}')
|
|
dist.init_process_group(
|
|
backend='nccl',
|
|
init_method='env://',
|
|
world_size=args.gpus,
|
|
rank=gpu)
|
|
|
|
# Create model.
|
|
model = resnet50(weights=None)
|
|
torch.cuda.set_device(gpu)
|
|
model.to(args.device)
|
|
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
|
model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
|
|
|
|
# Create dataloader.
|
|
train_dataloader = create_wds_dataloader(gpu, args, 'train')
|
|
eval_dataloader = create_wds_dataloader(gpu, args, 'eval')
|
|
|
|
# Optimizer.
|
|
optimizer = torch.optim.SGD(model.parameters(), 0.1)
|
|
|
|
# Main loop.
|
|
metric = torchmetrics.classification.Accuracy(top_k=1).to(args.device)
|
|
for epoch in range(1, args.epochs + 1):
|
|
if gpu == 0:
|
|
print(f'Running epoch {epoch}')
|
|
|
|
start = time.time()
|
|
train(model, args.device, train_dataloader, optimizer)
|
|
end = time.time()
|
|
if gpu == 0:
|
|
print(f'Training finished in {(end - start):>0.3f} seconds')
|
|
|
|
start = time.time()
|
|
evaluate(model, args.device, eval_dataloader, metric)
|
|
end = time.time()
|
|
if gpu == 0:
|
|
print(f'Evaluation finished in {(end - start):>0.3f} seconds')
|
|
|
|
if gpu == 0:
|
|
print('Done')
|
|
|
|
|
|
def create_args():
|
|
"""Create main args."""
|
|
parser = argparse.ArgumentParser(
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument(
|
|
'--gpus',
|
|
default=4,
|
|
type=int,
|
|
help='number of gpus to use')
|
|
parser.add_argument(
|
|
'--epochs',
|
|
default=1,
|
|
type=int,
|
|
help='number of total epochs to run')
|
|
parser.add_argument(
|
|
'--dataloader_num_workers',
|
|
default=2,
|
|
type=int,
|
|
help='number of workders for dataloader')
|
|
parser.add_argument(
|
|
'--train_data_path',
|
|
default='',
|
|
type=str,
|
|
help='path to training data')
|
|
parser.add_argument(
|
|
'--train_batch_size',
|
|
default=32,
|
|
type=int,
|
|
help='batch size for training per gpu')
|
|
parser.add_argument(
|
|
'--train_data_size',
|
|
default=50000,
|
|
type=int,
|
|
help='data size for training')
|
|
parser.add_argument(
|
|
'--eval_data_path',
|
|
default='',
|
|
type=str,
|
|
help='path to evaluation data')
|
|
parser.add_argument(
|
|
'--eval_batch_size',
|
|
default=32,
|
|
type=int,
|
|
help='batch size for evaluation per gpu')
|
|
parser.add_argument(
|
|
'--eval_data_size',
|
|
default=50000,
|
|
type=int,
|
|
help='data size for evaluation')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = create_args()
|
|
os.environ['MASTER_ADDR'] = 'localhost'
|
|
os.environ['MASTER_PORT'] = '8888'
|
|
|
|
args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
print(f'Launch job on {args.gpus} GPUs with DDP')
|
|
mp.spawn(worker, nprocs=args.gpus, args=(args,))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|