一、環境: 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