import torch
print(torch.cuda.is_available()) # 判断GPU是否可用
print(torch.cuda.device_count()) # 返回GPU个数
print(torch.cuda.get_device_name(0)) # 返回GPU名称，默认从0开始
print(torch.cuda.current_device()) # 返回当前设备索引还可以通过设备类型加上编号，来创建 device 对象：
device = torch.device('cuda', 0)
device = torch.device('cpu', 0)可以通过如下方式，设置当前 Python 脚本可见的 GPU.
- 在终端设置:
CUDA_VISIBLE_DEVICES=1 python my_script.py实例:
# Environment Variable Syntax Results
CUDA_VISIBLE_DEVICES=1 Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1 Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1" Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3 Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES="" No GPU will be visible- 在 Python 代码中设置
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"- 使用函数 set_device
import torch
torch.cuda.set_device(id)官方建议使用 CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。
默认使用CPU训练模型，使用GPU单机单卡训练时，模型和输入必须位于同一张GPU卡上。
使用GPU方法如下：
device = torch.device('cuda: 1')
model = model.cuda(device)
model = model.to(device) # 必须指定device需要用 torch.nn.DataParallel 来包装模型，包装模型时需要设置一些参数，包括：
参与训练的GPU有哪些， device_ids ;用于汇总梯度的GPU是哪个， output_device 。
nn.DataParallel 会自动将数据切分load到相应GPU，进行正向传播，并计算梯度，所以
batch_size应设置为 单个GPU上的batch_size * 使用的GPU个数 ：
import torch.nn as nn
model = MyNet()
model = model.cuda() # 需要这一步吗？
net = nn.DataParallel(
model, # 要被并行化的模块
device_ids=[0, 1, 2], # 整数列表或torch.device类型，默认值为所有的CUDA device
output_device=0 # 输出的device位置，默认值为device_ids[0]
)
output = net(input_var) # input_var可以在任何device上，包括CPU上。新的变动？注意：
是否需要手动将输入数据移动到GPU上？需要！是否需要 model.to(device)？需要！
由单进程控制多个GPU，所以在一个节点上只需要运行一次就可以创建一个进程。
上面的DataParallel进行数据并行时，使用一个进程来使用单个或多个GPU，若要创建多个进程就 需要代码并运行多次，很不方便。
有了distributed之后，我们可以只写一份代码，torch会自动将其分配给n个进程，分别在n个GPU 上运行。
pytorch为我们提供了torch.distributed.lanch启动器，用于在命令行分布式的执行Python 文件。可以这样获得当前进程（也即GPU的）的index:
parser = ArgumentParse()
parser.add_argument('--local_rank', default=-1, type=int,
help='node rank for distributed training')
args = parse.parse_args()
print(args.local_rank)然后使用 init_process_group 设置GPU通信所作用的后端和端口。初始化进程组主要有三种
方法：
- 显式指定
storerankword_size - 指定
init_method，来告诉进程如何去发现其他并行进程（init_method是个URL 字符串），并且指定rank和word_size或将它们写入URL字符串中。 - 如果以上两种都没有指定，则假设
init_method为env://
import torch.distributed as dist
dist.init_process_group(backend='nccl', # nccl的GPU支持目前是最好的，推荐使用
init_method='env://'
#world_size= ， # 整数，可选，任务的进程数。如果store给定，则必须指定
#rank= , # 当前进程的rank, 如果store给定，则必须指定
#store= #
)然后，使用DistributedSampler对数据集进行划分
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=...,
sampler=train_sampler)然后，使用 DistributedDataParallel 包装模型
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank])最后把数据和模型加载到当前进程使用的GPU中，进行向前向后传播：
torch.cuda.set_device(args.local_rank)
model.cuda()
for epoch in range(100):
for batch_idx, (data, target) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
target = target.cuda(non_blocking=True)
# ...
output = model(images)
loss = criterion(output, target)
# ...
optimizer.zero_grad()
loss.backward()
optimizer.step()在使用时，调用 torch.distributed.launch 启动器启动：
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py上面的launch使用过于繁琐，可以用更方便的 torch.multiprocessing 取代启动器。
下面的spawn 开启了 nprocs=4 个线程，每个线程执行 main_worker 并向其中传入
local_rank和args（即 4 和 myargs）作为参数：
import multiprocessing as mp
mp.spawn(main_worker, nprocs=4, args=(4,myargs))将原本需要 torch.distributed.launch 管理的执行内容，封装进 main_worker 函数中， 其中 proc 对应 local_rank（当前进程 index），ngpus_per_node 对应 4， args 对应 myargs：
def main_worker(proc, ngpus_per_node, args):
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
torch.cuda.set_device(args.local_rank)
# ...由于没有 torch.distributed.launch 读取的默认环境变量作为配置，我们需要手动为 init_process_group 指定参数：
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)在使用时，直接使用 python 运行就可以了：
python main.py其中backend为并行时GPU信息交换使用的后端，目前可选的有NCCL, Gloo, MPI. 对CPU和GPU的支持如下图：
目前Pytorch的 distributed 仅支持Linux平台。Gloo被包含在Pytorch安装文件中，NCCL则
在安装CUDA时被包含在内，而使用MPI则需要在用源码安装Pytorch时打开MPI支持。
如何选择后端：
- 分布式GPU训练：用NCCL。
- 分布式CPU训练：用Gloo。
- 用InfiniBand互连的GPU节点：用NCCL。只有NCCL支持IfiniBand和GPUDirect
- 用以太网互连的GPU节点：用NCCL，因为它的性能最佳。如果NCCL不能使用，再考虑Gloo。目前 Gloo后端比NCCL要慢。
- 用InfiniBand互连的CPU节点：如果InfiniBand可以使用IP，使用Gloo，否则使用MPI。 正在计划在未来版本中添加Gloo对InfiniBand的支持（时间点<2020-05-12 Tue>）。
- 用以太网互连的CPU节点：除非你有使用MPI的特殊理由，否则使用Gloo。
使用 init_process_group 设置GPU通信所作用的后端和端口。初始化进程组主要有三种方法：
- 显式指定
storerankword_size - 指定
init_method，来告诉进程如何去发现其他并行进程（init_method是个URL 字符串），并且指定rank和word_size或将它们写入URL字符串中。 - 如果以上两种都没有指定，则假设
init_method为env://
import torch.distributed as dist
dist.init_process_group(backend='nccl', # nccl的GPU支持目前是最好的，推荐使用
init_method='env://'
#world_size= ， # 整数，可选，任务的进程数。如果store给定，则必须指定
#rank= , # 当前进程的rank, 如果store给定，则必须指定
#store= #
)强烈建议用该方式来使用 DistributedDataParallel，使用多进程，每个进程使用一个 GPU。 这是目前 Pytorch 中，无论是单节点还是多节点，进行数据并行训练最快的方式。 在单节点多 GPU 上进行训练，该方式比 torch.nn.DataParallel 更快。这是因为分布是并行不需要 broadcast 参数。
假设每个主机有 N 个 GPUs，那么需要使用 N 个进程，并保证每个进程单独处理一个 GPU。 为了在每个主机（node）上使用多进程，可以使用 torch.distributed.launch 或 torch.multiprocessing.spawn 来实现。
在将数据移动到GPU上时，要将 non_blocking 设置为 False ，否则会导致不同GPU之间
不同步，从而卡进程，不能正常结束任务。
在 DataParallel 中，batch size 设置必须为单卡的 n 倍，但是在 DistributedDataParallel 内，batch size 设置于单卡一样即可。可见代码中设置的均为单个GPU所使用的batch size。
使用方法：
#on master node
python mnist-distributed.py -n 2 -g 4 --nr 0 --epoch 30
#on node 1
python mnist-distributed.py -n 2 -g 4 --nr 0 --epoch 30一般一个新的项目需要定义两个文件，一个 LightningModule 文件，一个 Trainer 文件。
这里有许多示例。
LightningModule 是 nn.Module 的严格超类。它提供了一个与 Trainer 交互的标准
界面。要开始自己的项目，可以从示例文件上修改，也可以自己定义 LightningModule 文件。
自定义的 LightningModule 文件要实现以下方法（参考这里）：
- 必需的方法：training_step, tng_dataloader, configure_optimizer
- 可选的方法：validation_step, validation_end, test_step, test_end, val_dataloader test_dataloader, on_save_checkpoint, on_load_checkpoint, update_tng_log_metrics, add_model_specific_args