NX入门教程软件篇-安装Pytorch（caffe2）带有cuda

NX入门教程软件篇-安装Pytorch（caffe2）带有cuda

说明：

• 介绍如何安装Pytorch（caffe2）
• python3.6 + Pytorch1.9.0 + cuda

步骤：

• 安装依赖

# Python 3.6 (download pip wheel from above)
$ sudo apt-get install libopenblas-base libopenmpi-dev python3-pip
$ pip3 install Cython

• 下载torch

$ wget https://nvidia.box.com/shared/static/2cqb9jwnncv9iii4u60vgd362x4rnkik.whl 
# 需要重名为正式的名称
$ mv 2cqb9jwnncv9iii4u60vgd362x4rnkik.whl  torch-1.9.0-cp36-cp36m-linux_aarch64.whl

测试：

• 新建脚本 vim hello_world.py
• 内容如下：

import torch
print(torch.cuda.is_available())
a = torch.cuda.FloatTensor(2)
print(a)
b = torch.randn(2).cuda()
print(b)
c = a + b
print(c)

• 运行

python3  hello_world.py

• 效果：

$ python3 hello_world.py 
True
tensor([0., 0.], device='cuda:0')
tensor([-1.4703, -0.2727], device='cuda:0')
tensor([-1.4703, -0.2727], device='cuda:0')

• 下载源码torchvision

$ git clone https://github.com/pytorch/vision
$ cd vision
$ sudo python3 setup.py install

测试：

• 新建脚本 vim hello_world2.py
• 内容如下：

import torch, torchvision
print("torch is: " + torch.__version__)
print("torchvision is : " + torchvision.__version__)

• 运行效果

$ python3 hello_world_vision.py 
torch is: 1.9.0
torchvision is : 0.11.0a0+ef71159