一、环境: Ubuntu16.04、 cuda8.0
二、方法:1.mkdir -p ~/dl/pytorch
2.下载pytorch,注意下载的版本,此次搭建成功的是V 0.1.10。下载地址:https://github.com/pytorch/pytorch
3.cd ~/dl/pytorch,将下载的pytorch解压后的文件放在该文件夹下。(4和5都在该文夹下)
4.sudo vim pytorch_jetson_install.sh,建立脚本文件。文件内容如下,然后./pytorch_jetson_install.sh运行该文件( 或sh pytorch_jetson_install.sh)
#!/bin/bash
#
# pyTorch install script for NVIDIA Jetson TX1/TX2,
# from a fresh flashing of JetPack 2.3.1 / JetPack 3.0 / JetPack 3.1
#
# for the full source, see jetson-reinforcement repo:
# https://github.com/dusty-nv/jetson-reinforcement/blob/master/CMakePreBuild.sh
#
# note: pyTorch documentation calls for use of Anaconda,
# however Anaconda isn't available for aarch64.
# Instead, we install directly from source using setup.py
sudo apt-get install python-pip
# upgrade pip
pip install -U pip
pip --version
# pip 9.0.1 from /home/ubuntu/.local/lib/python2.7/site-packages (python 2.7)
# clone pyTorch repo
cd pytorch
git init
git submodule update --init
# install prereqs
sudo pip install -U setuptools
sudo pip install -r requirements.txt
# Develop Mode:
python setup.py build_deps
sudo python setup.py develop
# Install Mode: (substitute for Develop Mode commands)
#sudo python setup.py install
# Verify CUDA (from python interactive terminal)
# import torch
# print(torch.__version__)
# print(torch.cuda.is_available())
# a = torch.cuda.FloatTensor(2)
# print(a)
# b = torch.randn(2).cuda()
# print(b)
# c = a + b
# print(c)
5.测试,vim test.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)
然后运行,python test.py
输出结果如下:
True
0
0
[torch.cuda.FloatTensor of size 2 (GPU 0)]
0.6851
-0.3392
[torch.cuda.FloatTensor of size 2 (GPU 0)]
0.6851
-0.3392
[torch.cuda.FloatTensor of size 2 (GPU 0)]
参考:
- https://gist.github.com/dusty-nv/ef2b372301c00c0a9d3203e42fd83426
- https://github.com/andrewadare/jetson-tx2-pytorch
- https://gist.github.com/dusty-nv/ef2b372301c00c0a9d3203e42fd83426