pytorch 分散式訓練初探

來自專欄 NeuronTalk神經雜談6 人贊了文章

pytorch 分散式訓練初探

為什麼需要分散式

眾所周知，深度神經網路發展到現階段，離不開GPU和數據。經過這麼多年的積累，GPU的計算能力越來越強，數據也積累的越來越多，大家會發現在現有的單機單卡或者單機多卡上很難高效地復現模型，甚至對於有些新的數據集來講，單機訓練簡直就是噩夢。

DatasetImagesMS COCO115,000Open Image dataset v41,740,000

為什麼單機8卡也會是噩夢呢？我們拿COCO和Google最近Release出來的Open Image dataset v4來做比較，訓練一個resnet152的檢測模型，在COCO上大概需要40個小時，而在OIDV4上大概需要40天，這還是在各種超參數正確的情況下，如果加上調試的時間，可能一個模型調完就該過年了吧。

所以這個時候我們需要分散式。

pytorch 分散式簡介

PyTorch分散式功能在0.4版本中得到大幅完善，它允許在多台機器之間交換Tensors，這樣就可以通過多台機器和更大的minibatch擴展網路訓練。

torch.distributed provides an MPI-like interface for exchanging tensor data across multi-machine networks. It supports a few different backends and initialization methods.

Currently torch.distributed supports four backends, each with different capabilities. The table below shows which functions are available for use with CPU / CUDA tensors. MPI supports cuda only if the implementation used to build PyTorch supports it.

pytorch文檔里有這麼一段介紹，但其實可選的只有Gloo。它是一個類似MPI的通信庫，你不需要考慮內存數據的拷貝，只需要實現邏輯就可以。

這裡引入了一個新的函數model = torch.nn.parallel.DistributedDataParallel(model)為的就是支持分散式模式

不同於原來在multiprocessing中的model = torch.nn.DataParallel(model,device_ids=[0,1,2,3]).cuda()函數，這個函數只是實現了在單機上的多GPU訓練，根據官方文檔的說法，甚至在單機多卡的模式下，新函數表現也會優於這箇舊函數。

這裡要提到兩個問題：

• 每個進程都有自己的Optimizer同時每個迭代中都進行完整的優化步驟，雖然這可能看起來是多餘的，但由於梯度已經聚集在一起並跨進程平均，因此對於每個進程都是相同的，這意味著不需要參數廣播步驟，從而減少了在節點之間傳輸張量tensor所花費的時間。
• 另外一個問題是Python解釋器的，每個進程都包含一個獨立的Python解釋器，消除了來自單個Python進程中的多個執行線程，模型副本或GPU的額外解釋器開銷和「GIL-thrashing」。 這對於大量使用Python運行時的模型尤其重要。

初始化

torch.distributed.init_process_group(backend, init_method=env://, **kwargs)

參數說明：

• backend(str): 後端選擇，包括上面那幾種 tcp mpi gloo
• init_method(str，optional): 用來初始化包的URL, 我理解是一個用來做並發控制的共享方式
• world_size(int, optional): 參與這個工作的進程數
• rank(int,optional): 當前進程的rank
• group_name(str,optional): 用來標記這組進程名的

解釋一下init_method()也有這三種方式，具體可參看http://pytorch.org/docs/master/distributed.html

• file:// 共享文件系統（要求所有進程可以訪問單個文件系統）有共享文件系統可以選擇
• tcp:// IP組播（要求所有進程都在同一個網路中）比較好理解，不過需要手動設置rank
• env:// 環境變數（需要您手動分配等級並知道所有進程可訪問節點的地址）默認是這個

分散式 Hello World

https://github.com/pytorch/examples/tree/master/imagenet

parser.add_argument(--world-size, default=2, type=int, help=number of distributed processes)parser.add_argument(--dist-url, default=tcp://172.16.1.186:2222, type=str, help=url used to set up distributed training)parser.add_argument(--dist-backend, default=gloo, type=str, help=distributed backend)parser.add_argument(--dist-rank, default=0, type=int, help=rank of distributed processes)

這幾個是必要的參數設置，其中最後一個是官網沒有的

if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size,rank=args.dist_rank)

這個是分散式的初始化，同樣，最後添加一個rank

model.cuda()model = torch.nn.parallel.DistributedDataParallel(model)

這裡，把我們平時使用的單機多卡，數據並行的API

model = torch.nn.DataParallel(model).cuda()

換掉即可。

if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

最後使用這個官方給的劃分方法，把數據集劃分即可。

需要注意的點

• 通信主機地址一定要寫ib的ip，否則帶寬不夠會造成Gloo底層觸發延遲錯誤
• 一定要加rank參數
• train_dataset最好不要用自己寫的sampler，否則還需要再實現一遍分散式的數據劃分方式

後續的計劃

一次講完分散式太難，所以這會是一個系列，系列結束後希望可以讓你自如地訓練自己的檢測、分割或者各種各樣的模型。

系列的內容包括但不限於:

• 使用uber/horovod加速訓練
• 詳解Train ImageNet In one Hour
• 如何得到更好的ResNet50
• 訓練檢測網路 or Mask RCNN

敬請期待~