Jetson Xavier NX配置pytorch环境

一、更换安装源

# 备份源文件，并打开
sudo cp /etc/apt/sources.list /etc/apt/sources.list.back
sudo gedit /etc/apt/sources.list

# 清华源
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-updates main restricted universe multiverse  
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-updates main restricted universe multiverse   
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-security main restricted universe multiverse   
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-security main restricted universe multiverse   
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-backports main restricted universe multiverse   
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial-backports main restricted universe multiverse   
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial main universe restricted   
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ xenial main universe restricted

# 中科大源
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-updates main restricted universe multiverse 
deb-src http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-updates main restricted universe multiverse 
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-security main restricted universe multiverse 
deb-src http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-security main restricted universe multiverse 
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-backports main restricted universe multiverse 
deb-src http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial-backports main restricted universe multiverse 
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial main universe restricted 
deb-src http://mirrors.ustc.edu.cn/ubuntu-ports/ xenial main universe restricted

# 如果签名验证错误
sudo apt-key adv --keyserver keyserver.ubuntu.com --recv-keys
sudo apt-get update

二、安装需要的依赖

sudo apt-get install libopenblas-base libopenmpi-dev
sudo pip3 install mpi4py
sudo pip3 install Cython

三、下载安装文件

1. 官网下载torch
   torch-1.5.0-cp36-cp36m-linux_aarch64(已是CUDA版本)

2. github下载torchvision
   选择torch对应的vision版本(torch1.5-->version0.6)

四、安装torch和vision

1. 安装torch

pip3 install numpy torch-1.1.0a0+b457266-cp36-cp36m-linux_aarch64.whl

2. 安装vision

# 解压
cd version
sudo python3 setup.py install

或者
sudo pip3 install  --no-deps torchvision==0.6.0

五、测试环境

# 进入python命令
python3

# 测试torch
import torch
print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
a = torch.cuda.FloatTensor(2).zero_()
print('Tensor a = ' + str(a))
b = torch.randn(2).cuda()
print('Tensor b = ' + str(b))
c = a + b
print('Tensor c = ' + str(c))

# 测试vision
import torchvision
print(torchvision.__version__)

参考