[AI Combat] DistributedDataParallel を使用して単一マシンのマルチカード並列トレーニングを実現する resnet50 モデル
DistributedDataParallel
pytorch フレームワークに基づく分散トレーニング ツール。
依存パッケージ
import argparse
import time
import torch
import torchvision
from torch import distributed as dist
from torchvision.models import resnet18
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn
事前トレーニング済みのモデルを読み込む
model_ft = models.resnet50(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层
使用DistributedDataParallel
net = model_ft
net.cuda()
net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = DDP(net, device_ids=[args.local_rank], output_device=args.local_rank)
カスタム データの読み込み
class MyDataset(torch.utils.data.Dataset):
def __init__(self, txt_path):
im_list = []
im_labels = []
with open(txt_path, 'r') as files:
for line in files:
#/x/y/a.jpg 1
#/x/y/b.jpg 2
items = line.split()
if len(items) != 2:
print(items)
continue
im_list.append(items[0])
im_labels.append(int(items[1]))
self.imgs = im_list
self.labels = im_labels
def __len__(self):
return len(self.imgs)
def __getitem__(self, item):
img_name = self.imgs[item]
label = self.labels[item]
def default_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
img = default_loader(img_name)
try:
img = data_tranforms(img)
except:
print("Cannot transform image: {}".format(img_name))
return img, label
data_tranforms = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
])
ダウンロードデータ:
data_root = 'dataset'
trainset = MyDataset(txt_path='./data/Train.txt')
valset = MyDataset(txt_path='./data/Test.txt')
sampler = DistributedSampler(trainset)
train_loader = DataLoader(trainset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
sampler=sampler)
val_loader = DataLoader(valset,
batch_size=batch_size,
shuffle=False,
pin_memory=True)
分散トレーニング
criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=lr)
net.train()
for e in range(epochs):
# DistributedSampler deterministically shuffle data
# by seting random seed be current number epoch
# so if do not call set_epoch when start of one epoch
# the order of shuffled data will be always same
sampler.set_epoch(e)
for idx, (imgs, labels) in enumerate(train_loader):
imgs = imgs.cuda()
labels = labels.cuda()
output = net(imgs)
loss = criterion(output, labels)
opt.zero_grad()
loss.backward()
opt.step()
reduce_loss(loss, global_rank, world_size)
if idx % 10 == 0 and global_rank == 0:
print('Epoch: {} step: {} loss: {}'.format(e, idx, loss.item()))
net.eval()
with torch.no_grad():
cnt = 0
total = len(val_loader.dataset)
for imgs, labels in val_loader:
imgs, labels = imgs.cuda(), labels.cuda()
output = net(imgs)
predict = torch.argmax(output, dim=1)
cnt += (predict == labels).sum().item()
if global_rank == 0:
print('eval accuracy: {}'.format(cnt / total))
モデルを保存
## 保存模型
path = './model/digit_classify-%s.pth' %(time.time())
torch.save(net.state_dict(), path)
print('*'*50)
print('data_tranforms', data_tranforms)
print('best model saved to ', path)
import shutil
shutil.copy(path, './model/digit_classify.pth')
完全なコード
import argparse
import time
import torch
import torchvision
from torch import distributed as dist
from torchvision.models import resnet18
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn
def reduce_loss(tensor, rank, world_size):
with torch.no_grad():
dist.reduce(tensor, dst=0)
if rank == 0:
tensor /= world_size
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, help="local gpu id")
args = parser.parse_args()
# python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=29500 train.py
batch_size = 96
world_size = 4
epochs = 20
lr = 0.001
NUMCLASS = 11
dist.init_process_group(backend='nccl', init_method='env://')
torch.cuda.set_device(args.local_rank)
global_rank = dist.get_rank()
print('global_rank', global_rank)
#net = resnet18()
model_ft = models.resnet18(pretrained=True)
#model_ft = models.resnet50(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层
#model_ft = model_ft.to(device)
net = model_ft
net.cuda()
net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = DDP(net, device_ids=[args.local_rank], output_device=args.local_rank)
data_tranforms = transforms.Compose([
transforms.Resize(448),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
])
class MyDataset(torch.utils.data.Dataset):
def __init__(self, txt_path):
im_list = []
im_labels = []
with open(txt_path, 'r') as files:
for line in files:
#/x/y/a.jpg 1
#/x/y/b.jpg 2
items = line.split()
if len(items) != 2:
print(items)
continue
im_list.append(items[0])
im_labels.append(int(items[1]))
self.imgs = im_list
self.labels = im_labels
def __len__(self):
return len(self.imgs)
def __getitem__(self, item):
img_name = self.imgs[item]
label = self.labels[item]
def default_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
img = default_loader(img_name)
try:
img = data_tranforms(img)
except:
print("Cannot transform image: {}".format(img_name))
return img, label
data_root = 'dataset'
trainset = MyDataset(txt_path='./data/Train.txt')
valset = MyDataset(txt_path='./data/Test.txt')
sampler = DistributedSampler(trainset)
train_loader = DataLoader(trainset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
sampler=sampler)
val_loader = DataLoader(valset,
batch_size=batch_size,
shuffle=False,
pin_memory=True)
criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=lr)
net.train()
for e in range(epochs):
# DistributedSampler deterministically shuffle data
# by seting random seed be current number epoch
# so if do not call set_epoch when start of one epoch
# the order of shuffled data will be always same
sampler.set_epoch(e)
for idx, (imgs, labels) in enumerate(train_loader):
imgs = imgs.cuda()
labels = labels.cuda()
output = net(imgs)
loss = criterion(output, labels)
opt.zero_grad()
loss.backward()
opt.step()
reduce_loss(loss, global_rank, world_size)
if idx % 10 == 0 and global_rank == 0:
print('Epoch: {} step: {} loss: {}'.format(e, idx, loss.item()))
net.eval()
with torch.no_grad():
cnt = 0
total = len(val_loader.dataset)
for imgs, labels in val_loader:
imgs, labels = imgs.cuda(), labels.cuda()
output = net(imgs)
predict = torch.argmax(output, dim=1)
cnt += (predict == labels).sum().item()
if global_rank == 0:
print('eval accuracy: {}'.format(cnt / total))
## 保存模型
path = './model/digit_classify-%s.pth' %(time.time())
torch.save(net.state_dict(), path)
print('*'*50)
print('data_tranforms', data_tranforms)
print('best model saved to ', path)
import shutil
shutil.copy(path, './model/digit_classify.pth')