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を使った深層学習の開発ができます.