PyTorch在Debian上支持多GPU和分布式并行計算,主要通過以下方式實現:
torch.nn.DataParallel
或torch.nn.parallel.DistributedDataParallel
(DDP)將模型復制到多個GPU,自動分配數據并行計算。
DataParallel
適合單機多卡,簡單易用。DDP
適合大規模分布式訓練(支持多機多卡),需配合torch.distributed.launch
或mp.spawn
啟動。apt-get install libnccl2 libnccl-dev
。CUDA_VISIBLE_DEVICES
控制可見GPU。python -m torch.distributed.launch --nproc_per_node=GPU數量 腳本.py
。init_method
指定主節點IP和端口。參考資料: