如何知道 pytorch 中使用了多少个 GPU?
how to know how many GPUs are used in pytorch?
我用来启动训练的 bash 文件如下所示:
CUDA_VISIBLE_DEVICES=3,4 python -m torch.distributed.launch \
--nproc_per_node=2 train.py \
--batch_size 6 \
--other_args
我发现每个GPU中张量的批量大小实际上是batch_size / num_of_gpu
= 6/2
= 3.
当我初始化我的网络时,我需要知道每个 GPU 中的批量大小。
(Ps。在这个阶段,我不能使用 input_tensor.shape
来获取批次维度的大小,因为 jet 中没有数据。)
不知何故我找不到pytorch在哪里存储参数--nproc_per_node
。
那么我怎么知道使用了多少 GPU,而不用手动将其传递为 --other_args
?
我认为您正在寻找 torch.distributed.get_world_size()
- 这会告诉您创建了多少个进程。
我用来启动训练的 bash 文件如下所示:
CUDA_VISIBLE_DEVICES=3,4 python -m torch.distributed.launch \
--nproc_per_node=2 train.py \
--batch_size 6 \
--other_args
我发现每个GPU中张量的批量大小实际上是batch_size / num_of_gpu
= 6/2
= 3.
当我初始化我的网络时,我需要知道每个 GPU 中的批量大小。
(Ps。在这个阶段,我不能使用 input_tensor.shape
来获取批次维度的大小,因为 jet 中没有数据。)
不知何故我找不到pytorch在哪里存储参数--nproc_per_node
。
那么我怎么知道使用了多少 GPU,而不用手动将其传递为 --other_args
?
我认为您正在寻找 torch.distributed.get_world_size()
- 这会告诉您创建了多少个进程。