Pytorch distributed 多卡並行載入模型操作

Posted on by

一、Pytorch distributed 多卡並行載入模型

這次來介紹下如何載入模型。

目前沒有找到官方的distribute 載入模型的方式，所以采用如下方式。

大部分情況下，我們在測試時不需要多卡並行計算。

所以，我在測試時隻使用單卡。

from collections import OrderedDict
device = torch.device("cuda")
model = DGCNN(args).to(device)  #自己的模型
state_dict = torch.load(args.model_path)    #存放模型的位置

new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
    # load params
model.load_state_dict (new_state_dict)

二、pytorch DistributedParallel進行單機多卡訓練

One_導入庫：

import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

Two_進程初始化：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=-1)
# 添加必要參數
# local_rank：系統自動賦予的進程編號，可以利用該編號控制打印輸出以及設置device

torch.distributed.init_process_group(backend="nccl", init_method='file://shared/sharedfile',
rank=local_rank, world_size=world_size)

# world_size：所創建的進程數，也就是所使用的GPU數量
# （初始化設置詳見參考文檔）

Three_數據分發：

dataset = datasets.ImageFolder(dataPath)
data_sampler = DistributedSampler(dataset, rank=local_rank, num_replicas=world_size)
# 使用DistributedSampler來為各個進程分發數據，其中num_replicas與world_size保持一致，用於將數據集等分成不重疊的數個子集

dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=1,drop_last=True, pin_memory=True, sampler=data_sampler)
# 在Dataloader中指定sampler時，其中的shuffle必須為False，而DistributedSampler中的shuffle項默認為True，因此訓練過程默認執行shuffle

Four_網絡模型：

torch.cuda.set_device(local_rank)
device = torch.device('cuda:'+f'{local_rank}')
# 設置每個進程對應的GPU設備

D = Model()
D = torch.nn.SyncBatchNorm.convert_sync_batchnorm(D).to(device)
# 由於在訓練過程中各卡的前向後向傳播均獨立進行，因此無法進行統一的批歸一化，如果想要將各卡的輸出統一進行批歸一化，需要將模型中的BN轉換成SyncBN
   
D = torch.nn.parallel.DistributedDataParallel(
D, find_unused_parameters=True, device_ids=[local_rank], output_device=local_rank)
# 如果有forward的返回值如果不在計算loss的計算圖裡，那麼需要find_unused_parameters=True，即返回值不進入backward去算grad，也不需要在不同進程之間進行通信。

Five_迭代：

data_sampler.set_epoch(epoch)
# 每個epoch需要為sampler設置當前epoch

Six_加載：

dist.barrier()
D.load_state_dict(torch.load('D.pth'), map_location=torch.device('cpu'))
dist.barrier()
# 加載模型前後用dist.barrier()來同步不同進程間的快慢

Seven_啟動：

CUDA_VISIBLE_DEVICES=1,3 python -m torch.distributed.launch --nproc_per_node=2 train.py --epochs 15000 --batchsize 10 --world_size 2
# 用-m torch.distributed.launch啟動，nproc_per_node為所使用的卡數，batchsize設置為每張卡各自的批大小

以上為個人經驗，希望能給大傢一個參考，也希望大傢多多支持WalkonNet。

推薦閱讀：