colab使用GPU

https://www.bilibili.com/video/BV1K3411876U
使用里面的代码

  1. ! pip install torch
  1. import torch
  2. torch.cuda.is_available()
  4. # 查看GPU数量
  5. torch.cuda.device_count()
  6. torch.cuda.get_device_name(0)
  8. #查看指定GPU容量
  9. torch.cuda.get_device_capability()
  11. #产看GPU的内存使用情况
  12. !nvidia-smi

下载DGL

有2种方式下载。一种是CPU版本的，一种是GPU版本的。
但是据我使用，反而是GPU版本的无法使用，会报错Device API gpu is not enabled. Please install the cuda version of dgl.可能是colab的问题。

CPU版本

  1. pip install dgl

然后使用官方的代码

  1. import dgl
  2. import torch as th
  3. u, v = th.tensor([0, 1, 2]), th.tensor([2, 3, 4])
  4. g = dgl.graph((u, v))
  5. g.ndata['x'] = th.randn(5, 3)   # 原始特征在CPU上
  6. g.device    
  7. #device(type='cpu')
  9. cuda_g = g.to('cuda:0')         # 这里会报错！！
  10. cuda_g.device
  11. #报错 Device API gpu is not enabled. Please install the cuda version of dgl.

其实，用另一种构造方法就可以了

  1. # 由GPU张量构造的图也在GPU上
  2. u, v = u.to('cuda:0'), v.to('cuda:0')
  3. g = dgl.graph((u, v))
  4. g.device
  5. device(type='cuda', index=0)

GPU版本

pip install dgl-cu111 -f https://data.dgl.ai/wheels/repo.html
这种方式我无法import DGL 。显示没有DGL这个module
colab里我的cuda是11.2。DGL官方文档里没有相应的11.2，所以我安装了11.1