PyTorchによるSSD Mobilenetでの転移学習（Jetson Nano）
Contents
転移学習
NVIDIAのJetson Nano 2GB 開発者キットで転移学習をやってみた時の備忘録。
PyTorchとOpenImages Dataset の画像を使って SSD-Mobilenet（mobilenet-v1-ssd-mp-0_675.pth） に新たに下記の８種類のフルーツの画像を学習させた。
- Apple（リンゴ）
- Orange（オレンジ）
- Banana（バナナ）
- Strawberry（イチゴ）
- Grape（ブドウ）
- Pear（洋梨）
- Pineapple（パイナップル）
- Watermelon（スイカ）
尚、学習手順はNVIDIAのHELLO AI WORLD の Re-training SSD-Mobilenet を参考にしている。
環境
今回の実行環境は以下の通り。
本体 | NVIDIA Jetson Nano 2GB 開発者キット |
OS | Ubuntu18.04.5 LTS |
JetPack SDK | JetPack 4.4.1 |
カメラ | ロジクール ウェブカメラ C270n |
パソコン | Windows10 Home Jetson Nano を ssh で操作する為のパソコン |
事前準備
OSのセットアップと初期設定、プロジェクトファイル類のコピーは終了している前提とする。
Jetson Nano 2GB 開発者キットのセットアップについては「Jetson Nano 2GB 開発者キットのセットアップ」の記事を参照して欲しい。
OSインストール後の環境設定については「NVIDIAのJetson Nanoで物体の検知をやってみた」のプロジェクトの準備の章を参照して欲しい。
デスクトップの無効化
この後の転移学習はマシンパワーをかなり使うのでメモリ等の消費量を抑えるために事前にデスクトップを無効化（command line interface）しておく。
sudo systemctl set-default multi-user.target尚、デスクトップに戻す時は以下のコマンドを実行する。
sudo systemctl set-default graphical.targetコマンド実行後に再起動を行うと設定が変更される。
フルーツ画像のダウンロード
Windows PowerShellの ssh で jetson Nano に接続して Dockerコンテナ を起動する。
※Jetson NanoのMicro USB Type-BとパソコンのUSBポートを接続すると192.168.55.1で接続できる。
ssh ユーザ名@192.168.55.1
cd jetson-inference/
docker/run.sh
該当のディレクトリーに移動して学習用のフルーツの画像を open_images_downloader.py スクリプトでダウンロードする。
cd /jetson-inference/python/training/detection/ssd
python3 open_images_downloader.py --class-names "Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon" --data=data/fruit–class-names で指定した種類のフルーツの画像が data/fruit ディレクトリにダウンロードされる。
ダウンロードは30～40Mbps（ダウンロード）の通信速度が出ている無線LAN環境で30分前後掛かった。
可能であれば有線LANに接続してから上記のコマンドを実行する事をオススメする。
約6,300枚の画像が data/fruit/ 配下に以下のディレクトリ構造でダウンロードされた。
├─jetson-inference
│ │
│ ├─python
│ │ │
│ │ ├──training
│ │ │ │
│ │ │ ├──detection
│ │ │ │ │
│ │ │ │ ├──ssd
│ │ │ │ │ │
│ │ │ │ │ ├──data
│ │ │ │ │ │ │
│ │ │ │ │ │ ├──fruit
│ │ │ │ │ │ │ │ csvファイル類（画像ファイルの情報）
│ │ │ │ │ │ │ │
│ │ │ │ │ │ │ ├──test
│ │ │ │ │ │ │ │ テスト用画像データ
│ │ │ │ │ │ │ │
│ │ │ │ │ │ │ ├──train
│ │ │ │ │ │ │ │ 学習用画像データ
│ │ │ │ │ │ │ │
│ │ │ │ │ │ │ └──validation
│ │ │ │ │ │ │ ハイパーパラメーター検証用画像データ
│ │ │ │ │ │ │転移学習
ダウンロードしたフルーツ画像を元にしてして転移学習を開始する。
尚、学習にはかなりのマシンリソースと時間が掛かり、途中でLAN接続が切れてしまう事もあったので ssh 接続では無く、直接キーボードとディスプレイを接続して以下のコマンドを実行した。
python3 train_ssd.py --data=data/fruit --model-dir=models/fruit --batch-size=4 --num-workers=0 --epochs=10パラメータの説明
train_ssd.py | SSD学習用プログラム（後述） |
–data | 学習用、テスト、パラメーター検証データの入っているディレクトリを指定する。 |
–model-dir | 学習済みモデルのチェックポイント（学習途中）の出力先ディレクトリーを指定する。
学習後に Epoch 0～9のモデルが出力される。 ファイル名の後半の5.314081・・・等の数値は損失関数の値を表しており、数値が少ないほど学習が進んでいることを示す。 尚、学習元のベースモデルは /models/ ディレクトリの mobilenet-v1-ssd-mp-0_675.pth を使用している。 |
–resume | トレーニングを再開するためのチェックポイントへのパスを指定する。 例えば先程のモデルの場合、Epoch=8 のモデルが一番損失関数が少なくて精度が良い。 –resume=models/fruit/mb1-ssd-Epoch-8-Loss-4.249728891470913.pth を指定するとそこから学習を再開する。 |
–batch-size | バッチサイズ（デフォルトは4） 学習途中でメモリ不足等でエラーになってしまう場合はバッチサイズを小さくすると記述がNVIDIAのページにあった。 |
–num-workers or –workers | 未指定の場合は2が指定されてマルチプロセス（2つ）でのデータの読み込みが行われる。 処理が高速になる反面メモリも消費するので初回の実行時は Epoch=2 の途中で “RuntimeError: DataLoader worker (pid 81) is killed by signal: Segmentation fault. ” で異常終了してしまった。 ※但し –workers=2 でも正常終了した時もあった
–num-workers=0 を指定してシングルプロセスでのデータ読み込みを行った。 |
–epochs | エポック数（学習回数） 10を指定した所、13時間41分で学習が終了した。 尚、推奨は100（但し学習時間が長くなる）とあった。 |
ソースコード
学習用 Python スクリプト（train_ssd.py）のソースコードは以下の通り。
jetson-inference のプロジェクトをコピーすると build/aarch64/bin/ ディレクトリーにある。
一部、日本語でコメントを追記した。
#
# train an SSD model on Pascal VOC or Open Images datasets
#
import os
import sys
import logging
import argparse
import itertools
import torch
from torch.utils.data import DataLoader, ConcatDataset
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
from vision.utils.misc import str2bool, Timer, freeze_net_layers, store_labels
from vision.ssd.ssd import MatchPrior
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.nn.multibox_loss import MultiboxLoss
from vision.ssd.config import vgg_ssd_config
from vision.ssd.config import mobilenetv1_ssd_config
from vision.ssd.config import squeezenet_ssd_config
from vision.ssd.data_preprocessing import TrainAugmentation, TestTransform
parser = argparse.ArgumentParser(
description='Single Shot MultiBox Detector Training With PyTorch')
# Params for datasets
parser.add_argument("--dataset-type", default="open_images", type=str,
help='Specify dataset type. Currently supports voc and open_images.')
parser.add_argument('--datasets', '--data', nargs='+', default=["data"], help='Dataset directory path')
parser.add_argument('--balance-data', action='store_true',
help="Balance training data by down-sampling more frequent labels.")
# Params for network
parser.add_argument('--net', default="mb1-ssd",
help="The network architecture, it can be mb1-ssd, mb1-lite-ssd, mb2-ssd-lite or vgg16-ssd.")
parser.add_argument('--freeze-base-net', action='store_true',
help="Freeze base net layers.")
parser.add_argument('--freeze-net', action='store_true',
help="Freeze all the layers except the prediction head.")
parser.add_argument('--mb2-width-mult', default=1.0, type=float,
help='Width Multiplifier for MobilenetV2')
# Params for loading pretrained basenet or checkpoints.
parser.add_argument('--base-net', help='Pretrained base model')
parser.add_argument('--pretrained-ssd', default='models/mobilenet-v1-ssd-mp-0_675.pth', type=str, help='Pre-trained base model')
parser.add_argument('--resume', default=None, type=str,
help='Checkpoint state_dict file to resume training from')
# Params for SGD
parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
help='Momentum value for optim')
parser.add_argument('--weight-decay', default=5e-4, type=float,
help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float,
help='Gamma update for SGD')
parser.add_argument('--base-net-lr', default=0.001, type=float,
help='initial learning rate for base net, or None to use --lr')
parser.add_argument('--extra-layers-lr', default=None, type=float,
help='initial learning rate for the layers not in base net and prediction heads.')
# Scheduler
parser.add_argument('--scheduler', default="cosine", type=str,
help="Scheduler for SGD. It can one of multi-step and cosine")
# Params for Multi-step Scheduler
parser.add_argument('--milestones', default="80,100", type=str,
help="milestones for MultiStepLR")
# Params for Cosine Annealing
parser.add_argument('--t-max', default=100, type=float,
help='T_max value for Cosine Annealing Scheduler.')
# Train params
parser.add_argument('--batch-size', default=4, type=int,
help='Batch size for training')
parser.add_argument('--num-epochs', '--epochs', default=30, type=int,
help='the number epochs')
parser.add_argument('--num-workers', '--workers', default=2, type=int,
help='Number of workers used in dataloading')
parser.add_argument('--validation-epochs', default=1, type=int,
help='the number epochs between running validation')
parser.add_argument('--debug-steps', default=10, type=int,
help='Set the debug log output frequency.')
parser.add_argument('--use-cuda', default=True, type=str2bool,
help='Use CUDA to train model')
parser.add_argument('--checkpoint-folder', '--model-dir', default='models/',
help='Directory for saving checkpoint models')
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
format='%(asctime)s - %(message)s', datefmt="%Y-%m-%d %H:%M:%S")
args = parser.parse_args()
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu") # cudaが利用可能ならGPUを使用する
if args.use_cuda and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True # cuDNN（GPUアクセラレーションライブラリ）が複数の畳み込みアルゴリズムをベンチマークして最速設定
logging.info("Using CUDA...")
def train(loader, net, criterion, optimizer, device, debug_steps=100, epoch=-1): # 学習
net.train(True)
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
for i, data in enumerate(loader): # index番号と要素（画像）数を取り出す
images, boxes, labels = data
images = images.to(device) # GPUに持っていく
boxes = boxes.to(device)
labels = labels.to(device)
optimizer.zero_grad() # Optimizer（最適化）勾配クリア
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes) # TODO CHANGE BOXES
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
if i and i % debug_steps == 0: # 100回に1回
avg_loss = running_loss / debug_steps
avg_reg_loss = running_regression_loss / debug_steps
avg_clf_loss = running_classification_loss / debug_steps
logging.info(
f"Epoch: {epoch}, Step: {i}/{len(loader)}, " +
f"Avg Loss: {avg_loss:.4f}, " +
f"Avg Regression Loss {avg_reg_loss:.4f}, " +
f"Avg Classification Loss: {avg_clf_loss:.4f}"
)
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
def test(loader, net, criterion, device): # 検証
net.eval()
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
num = 0
for _, data in enumerate(loader):
images, boxes, labels = data
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
num += 1
with torch.no_grad():
confidence, locations = net(images)
regression_loss, classification_loss = criterion(confidence, locations, labels, boxes)
loss = regression_loss + classification_loss
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
return running_loss / num, running_regression_loss / num, running_classification_loss / num
if __name__ == '__main__':
timer = Timer()
logging.info(args)
# make sure that the checkpoint output dir exists
if args.checkpoint_folder: # model-dirのディレクトリが無ければ作成する
args.checkpoint_folder = os.path.expanduser(args.checkpoint_folder)
if not os.path.exists(args.checkpoint_folder):
os.mkdir(args.checkpoint_folder)
# select the network architecture and config ネットワークの選択
if args.net == 'vgg16-ssd':
create_net = create_vgg_ssd
config = vgg_ssd_config
elif args.net == 'mb1-ssd': # デフォルト
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
elif args.net == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config
elif args.net == 'sq-ssd-lite':
create_net = create_squeezenet_ssd_lite
config = squeezenet_ssd_config
elif args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config
else:
logging.fatal("The net type is wrong.")
parser.print_help(sys.stderr)
sys.exit(1)
# create data transforms for train/test/val
train_transform = TrainAugmentation(config.image_size, config.image_mean, config.image_std)
target_transform = MatchPrior(config.priors, config.center_variance,
config.size_variance, 0.5)
test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)
# load datasets (could be multiple)
logging.info("Prepare training datasets.") # 学習用データセットの作成
datasets = []
for dataset_path in args.datasets: # --dataパラメーター
if args.dataset_type == 'voc': # Pascal VOC形式
dataset = VOCDataset(dataset_path, transform=train_transform,
target_transform=target_transform)
label_file = os.path.join(args.checkpoint_folder, "labels.txt")
store_labels(label_file, dataset.class_names)
num_classes = len(dataset.class_names)
elif args.dataset_type == 'open_images': # デフォルト
dataset = OpenImagesDataset(dataset_path,
transform=train_transform, target_transform=target_transform,
dataset_type="train", balance_data=args.balance_data)
label_file = os.path.join(args.checkpoint_folder, "labels.txt")
store_labels(label_file, dataset.class_names)
logging.info(dataset)
num_classes = len(dataset.class_names)
else:
raise ValueError(f"Dataset type {args.dataset_type} is not supported.")
datasets.append(dataset)
# create training dataset
logging.info(f"Stored labels into file {label_file}.")
train_dataset = ConcatDataset(datasets)
logging.info("Train dataset size: {}".format(len(train_dataset)))
train_loader = DataLoader(train_dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True)
# create validation dataset
logging.info("Prepare Validation datasets.") # 検証用データセット作成
if args.dataset_type == "voc":
val_dataset = VOCDataset(dataset_path, transform=test_transform,
target_transform=target_transform, is_test=True)
elif args.dataset_type == 'open_images':
val_dataset = OpenImagesDataset(dataset_path,
transform=test_transform, target_transform=target_transform,
dataset_type="test")
logging.info(val_dataset)
logging.info("Validation dataset size: {}".format(len(val_dataset)))
val_loader = DataLoader(val_dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=False)
# create the network
logging.info("Build network.")
net = create_net(num_classes)
min_loss = -10000.0
last_epoch = -1
# freeze certain layers (if requested)
base_net_lr = args.base_net_lr if args.base_net_lr is not None else args.lr
extra_layers_lr = args.extra_layers_lr if args.extra_layers_lr is not None else args.lr
if args.freeze_base_net:
logging.info("Freeze base net.")
freeze_net_layers(net.base_net)
params = itertools.chain(net.source_layer_add_ons.parameters(), net.extras.parameters(),
net.regression_headers.parameters(), net.classification_headers.parameters())
params = [
{'params': itertools.chain(
net.source_layer_add_ons.parameters(),
net.extras.parameters()
), 'lr': extra_layers_lr},
{'params': itertools.chain(
net.regression_headers.parameters(),
net.classification_headers.parameters()
)}
]
elif args.freeze_net:
freeze_net_layers(net.base_net)
freeze_net_layers(net.source_layer_add_ons)
freeze_net_layers(net.extras)
params = itertools.chain(net.regression_headers.parameters(), net.classification_headers.parameters())
logging.info("Freeze all the layers except prediction heads.")
else:
params = [
{'params': net.base_net.parameters(), 'lr': base_net_lr},
{'params': itertools.chain(
net.source_layer_add_ons.parameters(),
net.extras.parameters()
), 'lr': extra_layers_lr},
{'params': itertools.chain(
net.regression_headers.parameters(),
net.classification_headers.parameters()
)}
]
# load a previous model checkpoint (if requested)
timer.start("Load Model")
if args.resume: # resume指定されている時はチェックポイントから学習開始
logging.info(f"Resume from the model {args.resume}")
net.load(args.resume)
elif args.base_net: # base-netが指定されている時
logging.info(f"Init from base net {args.base_net}")
net.init_from_base_net(args.base_net)
elif args.pretrained_ssd: # デフォルトはmodels/mobilenet-v1-ssd-mp-0_675.pth
logging.info(f"Init from pretrained ssd {args.pretrained_ssd}")
net.init_from_pretrained_ssd(args.pretrained_ssd)
logging.info(f'Took {timer.end("Load Model"):.2f} seconds to load the model.')
# move the model to GPU
net.to(DEVICE)
# define loss function and optimizer
criterion = MultiboxLoss(config.priors, iou_threshold=0.5, neg_pos_ratio=3,
center_variance=0.1, size_variance=0.2, device=DEVICE)
optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
logging.info(f"Learning rate: {args.lr}, Base net learning rate: {base_net_lr}, "
+ f"Extra Layers learning rate: {extra_layers_lr}.")
# set learning rate policy
if args.scheduler == 'multi-step':
logging.info("Uses MultiStepLR scheduler.")
milestones = [int(v.strip()) for v in args.milestones.split(",")]
scheduler = MultiStepLR(optimizer, milestones=milestones,
gamma=0.1, last_epoch=last_epoch)
elif args.scheduler == 'cosine': # デフォルト
logging.info("Uses CosineAnnealingLR scheduler.")
scheduler = CosineAnnealingLR(optimizer, args.t_max, last_epoch=last_epoch)
else:
logging.fatal(f"Unsupported Scheduler: {args.scheduler}.")
parser.print_help(sys.stderr)
sys.exit(1)
# train for the desired number of epochs
logging.info(f"Start training from epoch {last_epoch + 1}.")
for epoch in range(last_epoch + 1, args.num_epochs): # Epochs分だけ繰り返す
scheduler.step()
train(train_loader, net, criterion, optimizer, # 学習
device=DEVICE, debug_steps=args.debug_steps, epoch=epoch)
if epoch % args.validation_epochs == 0 or epoch == args.num_epochs - 1:
val_loss, val_regression_loss, val_classification_loss = test(val_loader, net, criterion, DEVICE) # 検証
logging.info(
f"Epoch: {epoch}, " +
f"Validation Loss: {val_loss:.4f}, " +
f"Validation Regression Loss {val_regression_loss:.4f}, " +
f"Validation Classification Loss: {val_classification_loss:.4f}"
)
model_path = os.path.join(args.checkpoint_folder, f"{args.net}-Epoch-{epoch}-Loss-{val_loss}.pth")
net.save(model_path) # モデル保存
logging.info(f"Saved model {model_path}")
logging.info("Task done, exiting program.")モデルの変換
PyTorch 形式のモデル（*.pth）をNVIDIA TensorRT でも扱えるモデルの形式のONNX（Open Neural Network Exchange）に変換する。
TensorRT は NVIDIA製の GPU 向けに提供しているDeep Learning の推論を高速に実行するためのソフトウェア開発キット（SDK）だ。
また ONNX はオープンなフォーマットで複数のフレームワーク間（Keras、TensorFlowなど）でのモデルの互換が可能になっており、推論は PyTorch よりも高速化されている。
以下のプログラムを実行すると指定されたディレクトリからloss（損失）が一番少ないモデルを検出してONXXに変換する。
python3 onnx_export.py --model-dir=models/fruitパラメータの説明
onnx_export.py | 変換用プログラム |
–model-dir | 先程の学習でモデルを出力した models/fruit を指定した。 models/fruit 配下の複数のチェックポイントの中から loss（損失）が一番少ないモデルを自動で検出して、そのモデルをベースに ONXX 形式に変換する。 |
損失が一番少ない mb1-ssd-Epoch-3-Loss-4.175539761653786.pth を元に ssd-mobilenet.onnx が出力された。
ソースコード
onnx_export.py のソースコードは以下の通り。
NVIDIA が用意したスクリプトに一部日本語でコメントを追記している。
#
# converts a saved PyTorch model to ONNX format
#
import os
import sys
import argparse
import torch.onnx
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
# parse command line
parser = argparse.ArgumentParser()
parser.add_argument('--net', default="ssd-mobilenet", help="The network architecture, it can be mb1-ssd (aka ssd-mobilenet), mb1-lite-ssd, mb2-ssd-lite or vgg16-ssd.")
parser.add_argument('--input', type=str, default='', help="path to input PyTorch model (.pth checkpoint)")
parser.add_argument('--output', type=str, default='', help="desired path of converted ONNX model (default: <NET>.onnx)")
parser.add_argument('--labels', type=str, default='labels.txt', help="name of the class labels file")
parser.add_argument('--width', type=int, default=300, help="input width of the model to be exported (in pixels)")
parser.add_argument('--height', type=int, default=300, help="input height of the model to be exported (in pixels)")
parser.add_argument('--batch-size', type=int, default=1, help="batch size of the model to be exported (default=1)")
parser.add_argument('--model-dir', type=str, default='', help="directory to look for the input PyTorch model in, and export the converted ONNX model to (if --output doesn't specify a directory)")
args = parser.parse_args()
print(args)
# set the device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # cuDNN（GPUアクセラレーションライブラリ）が複数の畳み込みアルゴリズムをベンチマークして最速設定
print('running on device ' + str(device))
# format input model paths
if args.model_dir: # model-dirが指定されている時
args.model_dir = os.path.expanduser(args.model_dir) # フルPath取得
# find the checkpoint with the lowest loss
if not args.input:
best_loss = 10000
for file in os.listdir(args.model_dir): # model-dir配下のファイルを検索する
if not file.endswith(".pth"): # 拡張子がpthのファイルのみを対象にする
continue
try:
loss = float(file[file.rfind("-")+1:len(file)-4])
if loss < best_loss: # loss（損失関数）が一番少ないモデル（チェックポイント）を探す
best_loss = loss
args.input = os.path.join(args.model_dir, file)
except ValueError:
continue
print('found best checkpoint with loss {:f} ({:s})'.format(best_loss, args.input))
# append the model dir (if needed)
if not os.path.isfile(args.input):
args.input = os.path.join(args.model_dir, args.input)
if not os.path.isfile(args.labels):
args.labels = os.path.join(args.model_dir, args.labels)
# determine the number of classes
class_names = [name.strip() for name in open(args.labels).readlines()] # Labelからクラス名を取り出し
num_classes = len(class_names)
# construct the network architecture
print('creating network: ' + args.net)
print('num classes: ' + str(num_classes))
# ネットワークの指定
if args.net == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd' or args.net == 'ssd-mobilenet': # デフォルト
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif args.net == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif args.net == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
else:
print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
sys.exit(1)
# load the model checkpoint
print('loading checkpoint: ' + args.input)
net.load(args.input)
net.to(device) # GPU上で計算
net.eval() # モデルを推論モードにする
# create example image data
dummy_input = torch.randn(args.batch_size, 3, args.height, args.width).cuda() # 1（RGBかグレースケールか）×3（RGBの値）×300（X座標）×300（Y座標）のテンソルを作成する
# format output model path
if not args.output: # outputの指定が無ければ拡張子をonnxへ
args.output = args.net + '.onnx'
if args.model_dir and args.output.find('/') == -1 and args.output.find('\\') == -1:
args.output = os.path.join(args.model_dir, args.output)
# export to ONNX
input_names = ['input_0'] # 入力値と出力値に対する表示名（見やすさの為）
output_names = ['scores', 'boxes']
print('exporting model to ONNX...')
torch.onnx.export(net, dummy_input, args.output, verbose=True, input_names=input_names, output_names=output_names) # ONXXモデルのエクスポート
print('model exported to: {:s}'.format(args.output))
print('task done, exiting program')実行
実行の様子は以下の通り。
- Jetson Nano にロジクールのUSBカメラを接続する
- Jetson Nano にモニターをHDMI接続する（向かって左側）
- タブレット（HUAWEI MediaPad M5）に8種類のフルーツを表示させる（向かって右側）
- USBカメラに写ったフルーツの種類をリアルタイムで推論してバウンディングボックス（四角い枠）で囲む
リンク
リンク
リンク
プログラムの実行前に以下のコマンドでデスクトップ環境に戻しておく。
sudo systemctl set-default graphical.target 再起動でデスクトップ環境に戻る。
LXTerminal からDockerコンテナを起動して該当のディレクトリに移動する。
cd jetson-inference
docker/run.sh
cd python/training/detection/ssddetectnet コマンドで Mobile SSD を実行する。
detectnet --model=models/fruit/ssd-mobilenet.onnx --labels=models/fruit/labels.txt --input-blob=input_0 --output-cvg=scores --output-bbox=boxes /dev/video0パラメータの説明
–model | 先程作成したONNX形式のモデル（ssd-mobilenet.onnx）へのPathを指定する。 |
–labels | モデルのラベル（8種類のフルーツが登録されている）へのPathを指定する。 labels.txt |
–input-blob | 入力名：onnx_export.py で指定している input_0 |
–output-cvg | 出力名：onnx_export.py で指定している scores |
–output-bbox | 出力名：onnx_export.py で指定している boxes |
/dev/video0 | USBカメラのデバイスを指定した |
実行の様子
実行の様子は以下の動画を確認して欲しい。
8種類のフルーツを用意する事が出来なかったのでタブレットにフルーツを表示させて認識できるかどうかを試してみた。
イチゴ、リンゴ、オレンジ、ブドウ、洋梨は認識成功、パイナップルは画像によっては認識成功、バナナとスイカは認識しなかった。
本来であれば Epoch=30 程度は繰り返す所、10で終了しているのでloss（損失）が4%とそれほど精度がでていないのも理由の一つだと思う。
–resume を指定して、もう20 Epoch 程、学習をさせればもう少し精度が出るのだと思う。
最後に
以上で今回の記事は終了とする。
最後に
この記事が何処かで誰かの役に立つことを願っている。
尚、当記事中の商品へのリンクはAmazonアソシエイトへのリンクが含まれています。Amazonのアソシエイトとして、当メディアは適格販売により収入を得ていますのでご了承ください。
最近のコメント