目錄
pytorch多gpu並行訓練
注: 以下都在Ubuntu上面進行的調試, 使用的Ubuntu版本包括14, 18LST
參考文檔:
environment-variable-initialization
PYTORCH 1.0 DISTRIBUTED TRAINER WITH AMAZON AWS
pytorch/examples/imagenet/main.py
Getting Started with Distributed Data Parallel
1.單機多卡並行訓練
1.1.torch.nn.DataParallel
我一般在使用多GPU的時候, 會喜歡使用os.environ['CUDA_VISIBLE_DEVICES']來限制使用的GPU個數, 例如我要使用第0和第3編號的GPU, 那麼只需要在程序中設置:
os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
但是要注意的是, 這個參數的設定要保證在模型加載到gpu上之前, 我一般都是在程序開始的時候就設定好這個參數, 之後如何將模型加載到多GPU上面呢?
如果是模型, 那麼需要執行下面的這幾句代碼:
model = nn.DataParallel(model)
model = model.cuda()
如果是數據, 那麼直接執行下面這幾句代碼就可以了:
inputs = inputs.cuda()
labels = labels.cuda()
其實如果看pytorch官網給的示例代碼,我們可以看到下面這樣的代碼
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
model.to(device)
這個和我上面寫的好像有點不太一樣, 但是如果看一下DataParallel的內部代碼, 我們就可以發現, 其實是一樣的:
class DataParallel(Module):
def __init__(self, module, device_ids=None, output_device=None, dim=0):
super(DataParallel, self).__init__()
if not torch.cuda.is_available():
self.module = module
self.device_ids = []
return
if device_ids is None:
device_ids = list(range(torch.cuda.device_count()))
if output_device is None:
output_device = device_ids[0]
我截取了其中一部分代碼, 我們可以看到如果我們不設定好要使用的device_ids的話, 程序會自動找到這個機器上面可以用的所有的顯卡, 然後用於訓練. 但是因爲我們前面使用os.environ['CUDA_VISIBLE_DEVICES']限定了這個程序可以使用的顯卡, 所以這個地方程序如果自己獲取的話, 獲取到的其實就是我們上面設定的那幾個顯卡.
我沒有進行深入得到考究, 但是我感覺使用os.environ['CUDA_VISIBLE_DEVICES']對可以使用的顯卡進行限定之後, 顯卡的實際編號和程序看到的編號應該是不一樣的, 例如上面我們設定的是os.environ['CUDA_VISIBLE_DEVICES']="0,2", 但是程序看到的顯卡編號應該被改成了'0,1', 也就是說程序所使用的顯卡編號實際上是經過了一次映射之後纔會映射到真正的顯卡編號上面的, 例如這裏的程序看到的1對應實際的2
1.2.torch.nn.parallel.DistributedDataParallel
pytorch的官網建議使用DistributedDataParallel來代替DataParallel, 據說是因爲DistributedDataParallel比DataParallel運行的更快, 然後顯存分屏的更加均衡. 而且DistributedDataParallel功能更加強悍, 例如分佈式的模型(一個模型太大, 以至於無法放到一個GPU上運行, 需要分開到多個GPU上面執行). 只有DistributedDataParallel支持分佈式的模型像單機模型那樣可以進行多機多卡的運算.當然具體的怎麼個情況, 建議看官方文檔.
依舊是先設定好os.environ['CUDA_VISIBLE_DEVICES'], 然後再進行下面的步驟.
因爲DistributedDataParallel是支持多機多卡的, 所以這個需要先初始化一下, 如下面的代碼:
torch.distributed.init_process_group(backend='nccl', init_method='tcp://localhost:23456', rank=0, world_size=1)
第一個參數是pytorch支持的通訊後端, 後面會繼續介紹, 但是這裏單機多卡, 這個就是走走過場. 第二個參數是各個機器之間通訊的方式, 後面會介紹, 這裏是單機多卡, 設置成localhost就行了, 後面的端口自己找一個空着沒用的就行了. rank是標識主機和從機的, 這裏就一個主機, 設置成0就行了. world_size是標識使用幾個主機, 這裏就一個主機, 設置成1就行了, 設置多了代碼不允許.
其實如果是使用單機多卡的情況下, 根據pytorch的官方代碼distributeddataparallel, 是直接可以使用下面的代碼的:
torch.distributed.init_process_group(backend="nccl")
model = DistributedDataParallel(model) # device_ids will include all GPU devices by default
但是這裏需要注意的是, 如果使用這句代碼, 直接在pycharm或者別的編輯器中,是沒法正常運行的, 因爲這個需要在shell的命令行中運行, 如果想要正確執行這段代碼, 假設這段代碼的名字是main.py, 可以使用如下的方法進行(參考1 參考2):
python -m torch.distributed.launch main.py
注: 這裏如果使用了argparse, 一定要在參數裏面加上--local_rank, 否則運行還是會出錯的
之後就和使用DataParallel很類似了.
model = model.cuda()
model = nn.parallel.DistributedDataParallel(model)
但是注意這裏要先將model加載到GPU, 然後才能使用DistributedDataParallel進行分發, 之後的使用和DataParallel就基本一樣了
2.多機多gpu訓練
在單機多gpu可以滿足的情況下, 絕對不建議使用多機多gpu進行訓練, 我經過測試, 發現多臺機器之間傳輸數據的時間非常慢, 主要是因爲我測試的機器可能只是千兆網卡, 再加上別的一些損耗, 網絡的傳輸速度跟不上, 導致訓練速度實際很慢. 我看一個github上面的人說在單機8顯卡可以滿足的情況下, 最好不要進行多機多卡訓練
建議看這兩份代碼, 實際運行一下, 纔會真的理解怎麼使用
pytorch/examples/imagenet/main.py
2.1.初始化
初始化操作一般在程序剛開始的時候進行
在進行多機多gpu進行訓練的時候, 需要先使用torch.distributed.init_process_group()進行初始化. torch.distributed.init_process_group()包含四個常用的參數
backend: 後端, 實際上是多個機器之間交換數據的協議
init_method: 機器之間交換數據, 需要指定一個主節點, 而這個參數就是指定主節點的
world_size: 介紹都是說是進程, 實際就是機器的個數, 例如兩臺機器一起訓練的話, world_size就設置爲2
rank: 區分主節點和從節點的, 主節點爲0, 剩餘的爲了1-(N-1), N爲要使用的機器的數量, 也就是world_size
2.1.1.初始化backend
首先要初始化的是backend, 也就是俗稱的後端, 在pytorch的官方教程中提供了以下這些後端
[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片保存下來直接上傳(img-LhZ2L1aY-1571025664154)(img/2019-10-9-pytorch多gpu並行訓練/pytorch後端.png)]
根據官網的介紹, 如果是使用cpu的分佈式計算, 建議使用gloo, 因爲表中可以看到
gloo對cpu的支持是最好的, 然後如果使用gpu進行分佈式計算, 建議使用nccl, 實際測試中我也感覺到, 當使用gpu的時候, nccl的效率是高於gloo的. 根據博客和官網的態度, 好像都不怎麼推薦在多gpu的時候使用mpi
對於後端選擇好了之後, 我們需要設置一下網絡接口, 因爲多個主機之間肯定是使用網絡進行交換, 那肯定就涉及到ip之類的, 對於nccl和gloo一般會自己尋找網絡接口, 但是某些時候, 比如我測試用的服務器, 不知道是系統有點古老, 還是網卡比較多, 需要自己手動設置. 設置的方法也比較簡單, 在Python的代碼中, 使用下面的代碼進行設置就行:
import os
# 以下二選一, 第一個是使用gloo後端需要設置的, 第二個是使用nccl需要設置的
os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
我們怎麼知道自己的網絡接口呢, 打開命令行, 然後輸入ifconfig, 然後找到那個帶自己ip地址的就是了, 我見過的一般就是em0, eth0, esp2s0之類的, 當然具體的根據你自己的填寫. 如果沒裝ifconfig, 輸入命令會報錯, 但是根據報錯提示安裝一個就行了.
2.1.2.初始化init_method
初始化init_method的方法有兩種, 一種是使用TCP進行初始化, 另外一種是使用共享文件系統進行初始化
2.1.2.1.使用TCP初始化
看代碼:
import torch.distributed as dist
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
rank=rank, world_size=world_size)
注意這裏使用格式爲tcp://ip:端口號, 首先ip地址是你的主節點的ip地址, 也就是rank參數爲0的那個主機的ip地址, 然後再選擇一個空閒的端口號, 這樣就可以初始化init_method了.
2.1.2.2.使用共享文件系統初始化
好像看到有些人並不推薦這種方法, 因爲這個方法好像比TCP初始化要沒法, 搞不好和你硬盤的格式還有關係, 特別是window的硬盤格式和Ubuntu的還不一樣, 我沒有測試這個方法, 看代碼:
import torch.distributed as dist
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
rank=rank, world_size=world_size)
根據官網介紹, 要注意提供的共享文件一開始應該是不存在的, 但是這個方法又不會在自己執行結束刪除文件, 所以下次再進行初始化的時候, 需要手動刪除上次的文件, 所以比較麻煩, 而且官網給了一堆警告, 再次說明了這個方法不如TCP初始化的簡單.
2.1.3.初始化rank和world_size
這裏其實沒有多難, 你需要確保, 不同機器的rank值不同, 但是主機的rank必須爲0, 而且使用init_method的ip一定是rank爲0的主機, 其次world_size是你的主機數量, 你不能隨便設置這個數值, 你的參與訓練的主機數量達不到world_size的設置值時, 代碼是不會執行的.
2.1.4.初始化中一些需要注意的地方
首先是代碼的統一性, 所有的節點上面的代碼, 建議完全一樣, 不然有可能會出現一些問題, 其次, 這些初始化的參數強烈建議通過argparse模塊(命令行參數的形式)輸入, 不建議寫死在代碼中, 也不建議使用pycharm之類的IDE進行代碼的運行, 強烈建議使用命令行直接運行.
其次是運行代碼的命令方面的問題, 例如使用下面的命令運行代碼distributed.py:
python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 0 -ws 2
上面的代碼是在主節點上運行, 所以設置rank爲0, 同時設置了使用兩個主機, 在從節點運行的時候, 輸入的代碼是下面這樣:
python distributed.py -bk nccl -im tcp://10.10.10.1:12345 -rn 1 -ws 2
一定要注意的是, 只能修改rank的值, 其他的值一律不得修改, 否則程序就卡死了初始化到這裏也就結束了.
2.2.數據的處理-DataLoader
其實數據的處理和正常的代碼的數據處理非常類似, 但是因爲多機多卡涉及到了效率問題, 所以這裏纔會使用torch.utils.data.distributed.DistributedSampler來規避數據傳輸的問題. 首先看下面的代碼:
print("Initialize Dataloaders...")
# Define the transform for the data. Notice, we must resize to 224x224 with this dataset and model.
transform = transforms.Compose(
[transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# Initialize Datasets. STL10 will automatically download if not present
trainset = datasets.STL10(root='./data', split='train', download=True, transform=transform)
valset = datasets.STL10(root='./data', split='test', download=True, transform=transform)
# Create DistributedSampler to handle distributing the dataset across nodes when training
# This can only be called after torch.distributed.init_process_group is called
# 這一句就是和平時使用有點不一樣的地方
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
# Create the Dataloaders to feed data to the training and validation steps
train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=(train_sampler is None), num_workers=workers, pin_memory=False, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=False)
其實單獨看這段代碼, 和平時寫的很類似, 唯一不一樣的其實就是這裏先將trainset送到了DistributedSampler中創造了一個train_sampler, 然後在構造train_loader的時候, 參數中傳入了一個sampler=train_sampler. 使用這些的意圖是, 讓不同節點的機器加載自己本地的數據進行訓練, 也就是說進行多機多卡訓練的時候, 不再是從主節點分發數據到各個從節點, 而是各個從節點自己從自己的硬盤上讀取數據.
當然了, 如果直接讓各個節點自己讀取自己的數據, 特別是在訓練的時候經常是要打亂數據集進行訓練的, 這樣就會導致不同的節點加載的數據混亂, 所以這個時候使用DistributedSampler來創造一個sampler提供給DataLoader, sampler的作用自定義一個數據的編號, 然後讓DataLoader按照這個編號來提取數據放入到模型中訓練, 其中sampler參數和shuffle參數不能同時指定, 如果這個時候還想要可以隨機的輸入數據, 我們可以在DistributedSampler中指定shuffle參數, 具體的可以參考官網的api, 拉到最後就是DistributedSampler
2.3.模型的處理
模型的處理其實和上面的單機多卡沒有多大區別, 還是下面的代碼, 但是注意要提前想把模型加載到gpu, 然後纔可以加載到DistributedDataParallel
model = model.cuda()
model = nn.parallel.DistributedDataParallel(model)
2.4.模型的保存與加載
這裏引用pytorch官方教程的一段代碼:
def demo_checkpoint(rank, world_size):
setup(rank, world_size)
# setup devices for this process, rank 1 uses GPUs [0, 1, 2, 3] and
# rank 2 uses GPUs [4, 5, 6, 7].
n = torch.cuda.device_count() // world_size
device_ids = list(range(rank * n, (rank + 1) * n))
model = ToyModel().to(device_ids[0])
# output_device defaults to device_ids[0]
ddp_model = DDP(model, device_ids=device_ids)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"
if rank == 0:
# All processes should see same parameters as they all start from same
# random parameters and gradients are synchronized in backward passes.
# Therefore, saving it in one process is sufficient.
torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
# Use a barrier() to make sure that process 1 loads the model after process
# 0 saves it.
dist.barrier()
# configure map_location properly
rank0_devices = [x - rank * len(device_ids) for x in device_ids]
device_pairs = zip(rank0_devices, device_ids)
map_location = {'cuda:%d' % x: 'cuda:%d' % y for x, y in device_pairs}
ddp_model.load_state_dict(
torch.load(CHECKPOINT_PATH, map_location=map_location))
optimizer.zero_grad()
outputs = ddp_model(torch.randn(20, 10))
labels = torch.randn(20, 5).to(device_ids[0])
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
optimizer.step()
# Use a barrier() to make sure that all processes have finished reading the
# checkpoint
dist.barrier()
if rank == 0:
os.remove(CHECKPOINT_PATH)
cleanup()
我並沒有實際操作, 因爲多卡多GPU代碼運行起來實在是難受, 一次實驗可能就得好幾分鐘, 要是搞錯一點可能就得好幾十分鐘都跑不起來, 最主要的是還要等能用的GPU. 不過看上面的代碼, 最重要的實際是這句 dist.barrier(), 這個是來自torch.distributed.barrier(), 根據pytorch的官網的介紹, 這個函數的功能是同步所有的進程, 直到整組(也就是所有節點的所有GPU)到達這個函數的時候, 纔會執行後面的代碼, 看上面的代碼, 可以看到, 在保存模型的時候, 是隻找rank爲0的點保存模型, 然後在加載模型的時候, 首先得讓所有的節點同步一下, 然後給所有的節點加載上模型, 然後在進行下一步的時候, 還要同步一下, 保證所有的節點都讀完了模型. 雖然我不清楚這樣做的意義是什麼, 但是官網說不這樣做會導致一些問題, 我並沒有實際操作, 不發表意見.
至於保存模型的時候, 是保存哪些節點上面的模型, pytorch推薦的是rank=0的節點, 然後我看在論壇上, 有人也會保存所有節點的模型, 然後進行計算, 至於保存哪些, 我並沒有做實驗, 所以並不清楚到底哪種最好.