如何知道Pytork中使用了多少GPU？

Question

我用于启动培训的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中的批大小。 （另外，在这个阶段，我不能使用input_tensor.shape来获得批维度的大小，因为jet中没有数据馈送。）

不知何故，我找不到pytorch在哪里存储参数--nproc_per_node。 因此，我如何知道使用了多少GPU，而不将其作为--other_args手动传递

Answer 1

我认为您正在寻找^{}-这将告诉您创建了多少个进程