GPU训练

方式1：找到三个变量调用.cuda()

  1. 模型
  2. loss
  3. 数据 imgs, target = data

image.png

  1. mymodule = Mymodule()
  2. # ========================cuda======================================
  3. if torch.cuda.is_available():
  4.     mymodule = mymodule.cuda()
  6. # 3 损失函数和优化器
  7. loss_fn = nn.CrossEntropyLoss()
  8. # ========================cuda======================================
  9. if torch.cuda.is_available():
  10.     loss_fn = loss_fn.cuda()
  13.     with torch.no_grad():
  14.         for data in test_dataloader:
  15.             imgs, target = data
  16.             # ========================cuda======================================
  17.             if torch.cuda.is_available():
  18.                 imgs = imgs.cuda()
  19.                 target = target.cuda()

方式2：to(device)

  1. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

image.png

  1. # 定义训练的设备!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  2. # device = torch.device("cpu")
  3. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  4. print("设备：", device)
  7. mymodule = Mymodule()
  8. # ========================cuda======================================
  9. # if torch.cuda.is_available():
  10. #     mymodule = mymodule.cuda()
  11. mymodule = mymodule.to(device) # 也可以不用赋值 直接mymodule.to(device)
  14. # 3 损失函数和优化器
  15. loss_fn = nn.CrossEntropyLoss()
  16. # ========================cuda======================================
  17. # if torch.cuda.is_available():
  18. #     loss_fn = loss_fn.cuda()
  19. loss_fn = loss_fn.to(device) # 也可以不用赋值 直接loss_fn.to(device)
  22.     # 训练步骤开始
  23.     mymodule.train() # 非必要
  24.     for data in train_dataloader:
  25.         imgs, target = data
  26.         # # ========================cuda======================================
  27.         # if torch.cuda.is_available():
  28.         #     imgs = imgs.cuda()
  29.         #     target = target.cuda()
  30.         imgs = imgs.to(device)  # 必须这种
  31.         target = target.to(device) # 必须这种

配置多GPU

  2. import os
  3. # 配置gpu
  4. # os.environ["CUDA_VISIBLE_DEVICES"] =  "5,6,7"
  5. # os.environ["CUDA_VISIBLE_DEVICES"] =  "0"

torch.distributed.lanuch (shell脚本运行)

  3. #!/usr/bin/env bash
  5. CONFIG=$1
  6. GPUS=$2
  7. PORT=${PORT:-29500}
  9. PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
  10. python -m torch.distributed.launch --nproc_per_node=$GPUS --master_port=$PORT \
  11. $(dirname "$0")/train.py $CONFIG --launcher pytorch ${@:3}
  12. # 命令 torch.distributed.launch  python -m torch.distributed.launch main.py

https://zhuanlan.zhihu.com/p/86441879

python -m torch.distributed.launch --nproc_per_node=4 train.py --a b --c d
  • 其中python -m torch.distributed.lanuch表示调用torch.distributed.lanuch.py文件来进行分布式训练
  • -m表示将后面的torch.distributed.lanuch当做模块加载。
  • —nproc_per_node通常与GPU数量保持一致。
  • train.py才是真正的训练文件,后面是相关参数。