TX2 环境:cuda8.0 cudnn6.0 Ubuntu16.04 参考:https://github.com/dusty-nv/jetson-reinforcement/ $ mkdir pytorch $ git clone http://github.com/dusty-nv/jetson-reinforcement $ cd jetson-reinforcement $ git submodule update --init $ mkdir build $ cd build $ cmake ../ $ make
依次运行以上命令,所需时间较长,要保证网络畅通,所需内存也不小,保证tx2空间够用
验证pytorch是否安装好:
$ python 然后依次输入以下命令,
>>> 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))
会出现以下结果:
Python 2.7.12 (default, Nov 12 2018, 14:36:10) [GCC 5.4.0 20160609] on linux2 Type "help", "copyright", "credits" or "license" for more information. >>> import torch >>> print(torch.__version__) 0.3.0b0+af3964a >>> print('CUDA available: ' + str(torch.cuda.is_available())) CUDA available: True >>> a = torch.cuda.FloatTensor(2).zero_() >>> print('Tensor a = ' + str(a)) Tensor a = 0 0 [torch.cuda.FloatTensor of size 2 (GPU 0)] >>> b = torch.randn(2).cuda() >>> print('Tensor b = ' + str(b)) Tensor b = -0.3896 -0.5981 [torch.cuda.FloatTensor of size 2 (GPU 0)] >>> c = a + b >>> print('Tensor c = ' + str(c)) Tensor c = -0.3896 -0.5981 [torch.cuda.FloatTensor of size 2 (GPU 0)]