mmclassificationのインストールと動作確認

概要

mmclassificationとはOpenMMLabプロジェクトの一つであるPytorchをベースとしたオープンソース画像分類ツールボックスです。この記事ではmmclassifcationのインストールから実際に学習、推論の方法を紹介します。

動作環境

python 3,8 Pytorch 1.5+ cudatoolkit 11.3 torchvision 0.11.0

インストール

ローカル環境に構築する場合

Step 1. openmimをインストールし、mimを使用してMMCVのインストールを行います。

pip install -U openmim
mim install mmcv-full

Step 2. Githubからmmclassificationリポジトリのクローンを行い、mmclassificationライブラリのインストールを行う。

git clone https://github.com/open-mmlab/mmclassification.git
cd mmclassification
pip install -v -e .

# "-v" means verbose, or more output
# "-e" means installing a project in editable mode,
# thus any local modifications made to the code will take effect without reinstallation.

Step 3. インストールの確認

import mmcls
print(mmcls.__version__)
# Example output: 0.23.0 or newer

学習と推論

学習済みモデルのダウンロード

mmclassificationは多くの事前学習済み画像分類モデルを提供しています。これら事前学習済みモデルを用いることで私たちは簡単にモデルを動かすことができます。今回はmmclassificationが提供している事前学習済みモデルのひとつのResnet50を使用して学習、推論を実際に行ってみます。

checkpointファイルの準備

Step 1. mmclassificationディレクトリ直下にcheckpointsを保存しておくためのディレクトリを作成します。

cd mmclassification
mkdir checkpoints

Step 2. mmclassificationが提供しているチェックポイントファイルのダウンロードを行う。

wget -P ./checkpoints https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth

学習

提供されているチェックポイントと構成ファイルを使用して、学習を行います。

import mmcls
import mmcv

from mmcv import Config
from mmcls.utils import auto_select_device
from mmcls.apis import set_random_seed
from mmcls.datasets import build_dataset
from mmcls.models import build_classifier
from mmcls.apis import train_model

checkpoint_file="resnet50_8xb32_in1k_20210831-ea4938fc.pth"

cfg = Config.fromfile("configs/resnet/resnet50_8xb32_in1k.py")

#ワークディレクトリの作成
mmcv.mkdir_or_exist(abspath(cfg.work_dir))

# 分類器のビルド
model = build_classifier(cfg.model)
model.init_weights()

# データセットのビルド
datasets = [build_dataset(cfg.data.train)]

# クラス名を格納
model.CLASSES = datasets[0].CLASSES

train_model(
    model,
    datasets,
    cfg,
    distributed=False,
    validate=True,
    timestamp=time.strftime("%Y%m%d_%H%M%S", time.localtime()),
    meta=dict())

推論

提供されている事前学習ファイルを使用し、画像の推論を行いいます。

import mmcv
from mmcls.apis import inference_model, init_model

device = "cuda:0"
# device = "cpu"
checkpoint_file="resnet50_8xb32_in1k_20210831-ea4938fc.pth"
cfg　= "configs/resnet/resnet50_8xb32_in1k.py"
# モデルのビルド
model = init_model(cfg, checkpoint_file, device=device)

## 推論の実行
img = "./sample.jpg"
img_array =  mmcv.imread(img)
result.append(inference_model(model, img_array))

まとめ

今回はmmclassificationのインストール、学習、推論までを実際に行ってみました。また、構成ファイルを編集することで自作のデータセットを使用した学習やモデルのファインチューニングなどを行うことが可能になっています。

参考

mmclassification