介绍
WebDataset是一种大规模数据的组织格式，将大量数据以多个tar文件格式进行压缩，将随机文件读取过程变成一批文件的读取过程，大大提升了数据读取的速度。
本文介绍如何使用WebDataset 构建DataLoader并进行分布式训练。由于WebDataset 继承于Pytorch的iterable dataset，所以在构建DataLoader/分布式训练和常规的代码有所不同，但区别不会太大。
这里，我们展示用resnet50进行分布式训练imagenet，并且数据来源来自网络，不需要本地文件。
首先是一些必要的库
import os
import sys
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torchvision.models import resnet50
from torchvision import datasets, transforms
import ray
import webdataset as wds
import dataclasses
import time
from collections import deque
from typing import Optional
# Parameters
epochs = 5
maxsteps = int(100000)
数据集是非常容易的
bucket = "https://storage.googleapis.com/webdataset/fake-imagenet"
trainset_url = bucket + "/imagenet-train-{000000..000001}.tar"
batch_size = 8
cache_dir = "./_cache" # 把数据集保存到本地，如果不需要可以设为None
def make_dataloader_train():
# 构造一个函数make_sample，输入一个dict类型的sample，对其中的sample['jpg']进行处理，返回list，第一个元素是tensor图片（cpu)，第二个元素是一个int数据。
transform = transforms.Compose(
[
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
)
def make_sample(sample):
return transform(sample["jpg"]), sample["cls"]
# 首先获得url的列表，然后对url进行shuffle，增加数据的随机性；
# 读取tar文件并且保存tar文件到cache_dir中，注意resampled=True在分布式训练中必须使用！
trainset = wds.WebDataset(trainset_url, resampled=True, shardshuffle=True, cache_dir=cache_dir, nodesplitter=wds.split_by_node)
# 构造大小为1000的缓存空间，增加数据随机性，使用pil格式读取图片，调用make_sample返回数据
trainset = trainset.shuffle(1000).decode("pil").map(make_sample)
# 对于 IterableDataset 类型, batching 操作需要在数据集构造中完成
trainset = trainset.batched(64)
trainloader = wds.WebLoader(trainset, batch_size=None, num_workers=4)
# 对于resampled dataset，数据是无限的,可以设置一个epoch处理多少个batch的数据
trainloader = trainloader.with_epoch(100000)
return trainloader使用trainloader进行训练的代码和标准的Pytorch训练代码完全相同！
# 分布式训练的配置
@dataclasses.dataclass
class Config:
epochs: int = 10
max_steps: int = int(1e18)
lr: float = 0.001
momentum: float = 0.9
rank: Optional[int] = None
world_size: int = 2
backend: str = "nccl"
master_addr: str = "localhost"
master_port: str = "12355"
report_s: float = 15.0
report_growth: float = 1.1# 构造训练函数
# 每一个卡都要构造dataloader，建立model，损失函数和优化器，并且计算反向传播并进行梯度优化
def train(config):
# 构造模型
model = resnet50(pretrained=False).cuda()
if config.rank is not None:
model = DistributedDataParallel(model)
# 构造损失函数
loss_fn = nn.CrossEntropyLoss()
# 构造优化器
optimizer = torch.optim.SGD(model.parameters(), lr=config.lr)
# 构造trainloader
trainloader = make_dataloader()
# 训练循环和记录
losses, accuracies, steps = deque(maxlen=100), deque(maxlen=100), 0
for epoch in range(config.epochs):
for i, data in enumerate(trainloader):
# 在训练过程中单卡会把自己的数据放到gpu上
inputs, labels = data[0].cuda(), data[1].cuda()
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs)
# update statistics
loss = loss_fn(outputs, labels)
accuracy = (
(outputs.argmax(1) == labels).float().mean()
) # calculate accuracy
losses.append(loss.item())
accuracies.append(accuracy.item())
# verbose 可以修改成其他逻辑用来确定log的频率
if verbose and len(losses) > 0:
avgloss = sum(losses) / len(losses)
avgaccuracy = sum(accuracies) / len(accuracies)
print(
f"rank {config.rank} epoch {epoch:5d}/{i:9d} loss {avgloss:8.3f} acc {avgaccuracy:8.3f} {steps:9d}",
file=sys.stderr,
)
loss.backward()
optimizer.step()
steps += len(labels)
# 额外的保证总的数据量小于config.max_steps,可以进行快速调试
if steps > config.max_steps:
print(
"finished training (max_steps)",
steps,
config.max_steps,
file=sys.stderr,
)
return
print("finished Training", steps)
单卡快速调试代码
config = Config()
config.epochs = 1
config.max_steps = 1000
train(config)
多卡训练代码可以用torch.distributed.launch或者Ray等其他的，这些工具主要自动帮你设置MASTER_ADDR, MASTER_PORT和rank这些参数，做dist.init_process_group，但是最核心的还是会调用train(config)，比如下面的代码是Ray写的，逻辑是只调用一次distributed_training,然后这个函数调用train_on_ray去设施config.rank这些东西，然后最后就是每张卡会去自己训练。这里不同的工具可能有不同的逻辑，但是都会调用train(config)，config也不一定要用我这边的格式，config就是一种保存了rank这些需要用到的参数的数据结构。
Setting up Distributed Training with Ray
Ray is a convenient distributed computing framework. We are using it here to start up the training jobs on multiple GPUs. You can use torch.distributed.launch or other such tools as well with the above code. Ray has the advantage that it is runtime environment independent; you set up your Ray cluster in whatever way works for your environment, and afterwards, this code will run in it without change.
@ray.remote(num_gpus=1)
def train_on_ray(rank, config):
"""Set up distributed torch env and train the model on this node."""
# Set up distributed PyTorch.
if rank is not None:
os.environ["MASTER_ADDR"] = config.master_addr
os.environ["MASTER_PORT"] = config.master_port
dist.init_process_group(
backend=config.backend, rank=rank, world_size=config.world_size
)
config.rank = rank
# Ray will automatically set CUDA_VISIBLE_DEVICES for each task.
train(config)
if not ray.is_initialized():
ray.init()
ray.available_resources()["GPU"]
def distributed_training(config):
"""Perform distributed training with the given config."""
num_gpus = ray.available_resources()["GPU"]
config.world_size = min(config.world_size, num_gpus)
results = ray.get(
[train_on_ray.remote(i, config) for i in range(config.world_size)]
)
print(results)
config = Config()
config.epochs = epochs
config.max_steps = max_steps
config.batch_size = batch_size
print(config)
distributed_training(config)
这份代码里面的batched函数因为是操作tensor 和 int，所以有默认的batched方式，但是如果batched函数需要操作一些不需要batched操作的部分，比如我们的数据还有info，info的部分保存了图片路径，那么我们需要自己实现batched的函数，类似于实现make_sample函数。我们注意到我们的trainloader最后的输出是一个list结构，list[0]是shape为[B, H, W, C]的cpu上的tensor, list[1] 是shape为[B, 1]的cpu上面的tensor,这关系到training loop里面我们要如何从data构造inputs 和labels。
如果想要跑一下代码，请跳转到train-resnet50-multiray-wds demo。
总结：
为了让整个代码逻辑清晰，可以把代码分成若干个独立的函数，包括：
1. make_dataloader （需要考虑分布式训练的resampled=True和设置loader的with_epoch）
2. make_model （分布式训练需要DistributedDataParallel）
3. make_optimizer
4. make_loss
5. 组织1~4形成train(config)，里面要关注的就是data的数据用.cuda变成cuda的数据。
6. 根据你使用的训练框架写代码来调用train(config)。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
