在jittor中有类似torch.cuda可以查看显卡数量和显卡信息的函数吗？

问答
1

在python中import jittor后，看文档里似乎没有类似torch.cuda的函数

2

jittor 中可以通过 get_device_count() 获取显卡数量，通过 display_memory_info() 打印 cpu、gpu 的总内存、可用内存等信息。

>>> import jittor as jt
>>> jt.get_device_count()
4
>>> jt.display_memory_info()
[i 1206 17:41:29.071518 44 <stdin>:1]
=== display_memory_info ===
 total_cpu_ram: 62.78GB total_cuda_ram: 23.65GB
 hold_vars: 0 lived_vars: 0 lived_ops: 0
 update queue: 0/0
 name: sfrl is_cuda: 1 used:     0 B(-nan%) unused:     0 B(-nan%) total:     0 B
 name: sfrl is_cuda: 0 used:     0 B(-nan%) unused:     0 B(-nan%) total:     0 B
 name: sfrl is_cuda: 0 used:     0 B(-nan%) unused:     0 B(-nan%) total:     0 B
 cpu&gpu:     0 B gpu:     0 B cpu:     0 B
 free: cpu(11.23GB) gpu(12.22GB)
===========================
3
  • 那么在jt.flags.use_cuda = 1后，代码当前是跑在第一个gpu吗？
  • 在mpi里提到可以用CUDA_VISIBLE_DEVICES="2,3"选择gpu，在代码里可以指定将数据载入到某个gpu吗，比如.cuda(1)或者set_device？
4

jt.flags.use_cuda 代表是否使用 cuda，不控制在哪张卡运行。

如果您仅使用单个GPU，那么就使用环境变量 CUDA_VISIBLE_DEVICES=x 指定 gpu。此环境变量限制了接下来运行程序能够使用到的 GPU，jittor 会把数据载入到可用的第一块 gpu 上。Jittor 没有在 python 层面提供手动选择 gpu 的接口。

如果您希望使用多卡，除了指定 CUDA_VISIBLE_DEVICES="2,3" 外，还需要利用 jittor 的 mpi 的接口更改相关代码。