Ubuntu 18.04にCUDA9.0を入れてPyTorchを動かす

このドキュメントの内容

• ubuntu18.04, CUDA9.0の環境下に、torch==1.0.1を入れる手順のメモ.

前提

• OS: Ubuntu 18.04
• GPU: GeForce GTX 1080
• CUDA9.0:
  • この手順でインストールした.
• python3.7
• pip3
• pipenv

手順

python3.7の仮想環境を作る。

$pipenv install --python 3.7

pipenvでインストールを行う。
コマンドは公式サイト参照。

$ pipenv install torch==1.0.1

動作確認

$ pipenv shell
$ python3
Python 3.7.3 (default, Apr  3 2019, 19:16:38) 
[GCC 8.0.1 20180414 (experimental) [trunk revision 259383]] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> 

>>> import torch
>>> torch.cuda.current_device()
0
>>> torch.cuda.device(0)
<torch.cuda.device object at 0x7f7e55f3e128>
>>> torch.cuda.device_count()
1
>>> torch.cuda.get_device_name(0)
'GeForce GTX 1080'
>>> torch.cuda.is_available()
True

以上.