mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
235 lines
6.9 KiB
Python
235 lines
6.9 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 os
|
|
import time
|
|
|
|
from PIL import Image
|
|
import torch
|
|
from torch import nn
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
import torchmetrics
|
|
import torchvision
|
|
from torchvision.models import resnet50
|
|
|
|
|
|
class ImageFolder(torchvision.datasets.ImageFolder):
|
|
"""Class for loading imagenet."""
|
|
|
|
def __init__(self, image_list_file, transform=None, target_transform=None):
|
|
self.samples = self._make_dataset(image_list_file)
|
|
self.loader = self._loader
|
|
|
|
self.imgs = self.samples
|
|
self.targets = [s[1] for s in self.samples]
|
|
|
|
self.transform = transform
|
|
self.target_transform = target_transform
|
|
|
|
def _make_dataset(self, image_list_file):
|
|
items = []
|
|
with open(image_list_file, 'r') as f:
|
|
for line in f:
|
|
item = line.strip().split(' ')
|
|
items.append((item[0], int(item[1])))
|
|
return items
|
|
|
|
def _loader(self, image_path):
|
|
with open(image_path, 'rb') as f:
|
|
img = Image.open(f)
|
|
img = img.convert('RGB')
|
|
return img
|
|
|
|
|
|
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 train dataloader.
|
|
train_dataset = ImageFolder(
|
|
image_list_file=args.train_data_path,
|
|
transform=torchvision.transforms.Compose([
|
|
torchvision.transforms.RandomResizedCrop(224),
|
|
torchvision.transforms.RandomHorizontalFlip(),
|
|
torchvision.transforms.ToTensor(),
|
|
torchvision.transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]))
|
|
train_sampler = torch.utils.data.distributed.DistributedSampler(
|
|
train_dataset, num_replicas=args.gpus, rank=gpu)
|
|
train_dataloader = torch.utils.data.DataLoader(
|
|
dataset=train_dataset,
|
|
batch_size=args.train_batch_size,
|
|
shuffle=False,
|
|
num_workers=args.dataloader_num_workers,
|
|
pin_memory=True,
|
|
sampler=train_sampler)
|
|
if gpu == 0:
|
|
print(f'Train dataloader | samples: {len(train_dataloader.dataset)}, '
|
|
f'num workers: {train_dataloader.num_workers}, '
|
|
f'global batch size: {args.train_batch_size * args.gpus}, '
|
|
f'batches/epoch: {len(train_dataloader)}')
|
|
|
|
# Create eval dataloader.
|
|
eval_dataset = ImageFolder(
|
|
image_list_file=args.eval_data_path,
|
|
transform=torchvision.transforms.Compose([
|
|
torchvision.transforms.Resize(256),
|
|
torchvision.transforms.CenterCrop(224),
|
|
torchvision.transforms.ToTensor(),
|
|
torchvision.transforms.Normalize(
|
|
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
]))
|
|
eval_sampler = torch.utils.data.distributed.DistributedSampler(
|
|
eval_dataset, num_replicas=args.gpus, rank=gpu)
|
|
eval_dataloader = torch.utils.data.DataLoader(
|
|
dataset=eval_dataset,
|
|
batch_size=args.eval_batch_size,
|
|
shuffle=False,
|
|
num_workers=args.dataloader_num_workers,
|
|
pin_memory=True,
|
|
drop_last=True,
|
|
sampler=eval_sampler)
|
|
if gpu == 0:
|
|
print(f'Eval dataloader | samples: {len(eval_dataloader.dataset)}, '
|
|
f'num workers: {eval_dataloader.num_workers}, '
|
|
f'batch size: {args.eval_batch_size}, '
|
|
f'batches/epoch: {len(eval_dataloader)}')
|
|
|
|
# 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}')
|
|
train_sampler.set_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(
|
|
'--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')
|
|
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()
|