pytouch使用多GPU訓練
目錄
- 目錄
- pytorch多gpu並行訓練
- 1.單機多卡並行訓練
- 1.1.torch.nn.DataParallel
- 1.2.如何平衡DataParallel帶來的視訊記憶體使用不平衡的問題
- 1.3.torch.nn.parallel.DistributedDataParallel
- 2.多機多gpu訓練
- 2.1.初始化
- 2.1.1.初始化backend
- 2.1.2.初始化init_method
- 2.1.2.1.使用TCP初始化
- 2.1.2.2.使用共享檔案系統初始化
- 2.1.3.初始化rank和world_size
- 2.1.4.初始化中一些需要注意的地方
- 2.2.資料的處理-DataLoader
- 2.3.模型的處理
- 2.4.模型的儲存與載入
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.如何平衡DataParallel帶來的視訊記憶體使用不平衡的問題
這個問題其實討論的也比較多了, 官方給的解決方案就是使用DistributedDataParallel來代替DataParallel(實際上DistributedDataParallel視訊記憶體分配的也不是很平衡), 但是從某些角度來說,DataParallel使用起來確實比較方便, 而且最近使用DistributedDataParallel遇到一些小問題. 所以這裡提供一個解決視訊記憶體使用不平衡問題的方案:
首先這次的解決方案來自transformer-XL的官方程式碼:https://github.com/kimiyoung/transformer-xl
然後我將其中的平衡GPU視訊記憶體的程式碼提取了出來(原始碼好像有點小問題)放到了github上面:https://github.com/Link-Li/Balanced-DataParallel
這裡的程式碼是原作者繼承了DataParallel類之後進行了改寫:
class BalancedDataParallel(DataParallel):
def __init__(self, gpu0_bsz, *args, **kwargs):
self.gpu0_bsz = gpu0_bsz
super().__init__(*args, **kwargs)
...這個BalancedDataParallel類使用起來和DataParallel類似, 下面是一個示例程式碼:
my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()這裡包含三個引數, 第一個引數是第一個GPU要分配多大的batch_size, 但是要注意, 如果你使用了梯度累積, 那麼這裡傳入的是每次進行運算的實際batch_size大小. 舉個例子, 比如你在3個GPU上面跑程式碼, 但是一個GPU最大隻能跑3條資料, 但是因為0號GPU還要做一些資料的整合操作, 於是0號GPU只能跑2條資料, 這樣一算, 你可以跑的大小是2+3+3=8, 於是你可以設定下面的這樣的引數:
batch_szie = 8
gpu0_bsz = 2
acc_grad = 1
my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()這個時候突然想跑個batch size是16的怎麼辦呢, 那就是4+6+6=16了, 這樣設定累積梯度為2就行了:
batch_szie = 16
gpu0_bsz = 4
acc_grad = 2
my_net = MyNet()
my_net = BalancedDataParallel(gpu0_bsz // acc_grad, my_net, dim=0).cuda()1.3.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_size2.1.1.初始化backend
首先要初始化的是backend, 也就是俗稱的後端, 在pytorch的官方教程中提供了以下這些後端
根據官網的介紹, 如果是使用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的節點, 然後我看在論壇上, 有人也會儲存所有節點的模型, 然後進行計算, 至於儲存哪些, 我並沒有做實驗, 所以並不清楚到底哪種最好.