jetsonで深層学習をする環境を整える


はじめに

  • jetsonはarmアーキテクチャなので環境構築が面倒くさい
  • nvidia公式のjetson用深層学習イメージがある nvidia NGC
  • vscodeを使って快適に開発を行いたい

流れ

  • vscodeのインストール
  • vscodeにdocker extentionをインストール
  • dockerhubからimageをpullする
  • コンテナを立ち上げる
  • vscodeでコンテナにアタッチする

vscodeのインストール

ここからダウンロードしてインストールできます.

jetsonはarmアーキテクチャなのでARM64を選択してください

docker extentionのインストール

vscodeの拡張機能からdockerをインストールします.

imageをpullする

nvidia NGCからjetson用のimageをpullします.今回はpytorchを使いたいのでこれを使います.
jetpackのバージョンに合わせないといけないのでjetpackのバーションを確認します

> dpkg-query --show nvidia-l4t-core
nvidia-l4t-core 32.4.4-20201016123640

このバージョンにあったものをpullします

docker pull nvcr.io/nvidia/l4t-pytorch:r32.4.4-pth1.6-py3

コンテナを立ち上げる

ターミナルから

docker run --rm --runtime nvidia --network host -v /home/user/project:/location/in/container nvcr.io/nvidia/l4t-pytorch:r32.4.4-pth1.6-py3

と入力するとコンテナが立ち上がります
--rm はコンテナ終了時に自動的にコンテナを削除するコマンド (削除したくない場合は--rmコマンドをつけない)
--runtime nvidia これをつけるとjetsonのcuda環境をdockerに提供することができる
-v : 左がマウントしたいホストのファイルパス(絶対パスでないといけないらしい)。右がコンテナ内のファイルの置き場。

docker extentionからもコンテナが立ち上がっていることが確認できる.

vscodeでコンテナにアタッチする

vscodeのdocker extensionからcontainersの起動中のコンテナ上で右クリックしattach visual studio code を選択

vscodeで開発する環境が整いました

実際に動くかの確認

GPUが認識されているか確認しましょう

import torch
# GPUの確認
use_cuda = torch.cuda.is_available()
print('Use CUDA:', use_cuda)

実行結果は以下のようになります

Use CUDA: True

dockerコンテナでGPUを使った深層学習の開発ができます.