Google colabでKaolin使ってみた

Kaolinとは3Dディープラーニングの研究を加速することを目的としたPyTorchライブラリ。
つまり簡単に3Dモデルの機械学習ができる。

現在(2020/3/8時点)ではwindows環境はサポートされておらずubuntu等を使ってLinux環境を構築する必要がある。またGPUスペックの自信もないため、colabで試しにやってみた。
今回は\kaolin-master\examples\Classificationにある「point_cloud_classification.ipynb」の「イス」か「ソファ」かのクラス分類を行う。

Kaolinリンク

環境構築

・Kaolinインストール

Google Colaboratoryは基本的なライブラリはインストールされているが、Kaolinはまだまだ開発中なのでインストールされていない。
まずはcolabにKaolinをpipインストールする。
もちろん事前にランタイムのタイプをGPUにしておく。
※インストールに10分かかりました。学習の都度10分待つのはしんどいがしょうがない。

!pip install git+https://github.com/NVIDIAGameWorks/kaolin.git
import kaolin as kal
print(kal.__version__)

colabは最新のライブラリを取り揃えているため、3つのライブラリのバージョンが互換性がないようなのでランタイムを再起動させます。

WARNING: The following packages were previously imported in this runtime:
 [matplotlib,mpl_toolkits,sphinxcontrib]
You must restart the runtime in order to use newly installed versions.

・モジュールのインポート

モジュールも使うようなのでgithubからダウンロードし解凍済みkaolin-masterをgoogleドライブにアップロードする。
ドライブをマウントしてモジュールをインポートする。
パスは左側のフォルダマークを右クリックしてパスをコピーすれば取得できる。colabのパスは特殊なのでこれが確実。

import sys
ROOT_PATH = '/content/drive/My Drive/xxxxx/kaolin-master'
sys.path.append(ROOT_PATH)

・データセットのアップロード

Princeton ModelNet よりデータセットModelNet10をダウンロードする。10種類のデータセットがあるため今回使う「chair」と「sofa」をドライブにアップロードする。これらだけで1GBくらいある。
パスは必ず下記のような形式にしておく。
ModelNet10 ─ chair
                     └ sofa

データセット作成者から元の著作権を複製

**Copyright**
All CAD models are downloaded from the Internet and the original authors hold the copyright of the CAD models. The label of the data was obtained by us via Amazon Mechanical Turk service and it is provided freely. This dataset is provided for the convenience of academic research only.

データセットの中身はこんな感じ
chair : train 889ファイル
          test   100ファイル
sofa  : train 680ファイル
          test   100ファイル

拡張子は.offなのでMeshLabを使えば見れる。
chairのモデルを4つ見てみた(左下のやつがものすごくデカかった)

sofaのモデルを4つ

これで環境設定は完了。

点群のクラス分類

・必要なライブラリのインポート

import torch
from torch.utils.data import DataLoader
import kaolin as kal
from kaolin import ClassificationEngine
from kaolin.datasets import ModelNet
from kaolin.models.PointNet import PointNetClassifier as PointNet
import kaolin.transforms as tfs

・パラメータの定義

modelnet_path = '/content/drive/My Drive/xxxxx/ModelNet10'
categories = ['chair', 'sofa']
num_points = 1024
device = 'cuda'

"modelnet_path"　モデルのパス
"categories"　今回分類するクラス
"num_points"　メッシュからサンプリングする点群の数
"device"　GPUを使う指示(colabならいらないかもしれない)

・正規化

transform = tfs.Compose([
   tfs.TriangleMeshToPointCloud(num_samples=num_points),
   tfs.NormalizePointCloud()
])

メッシュデータから点群データに変換し、標準偏差1で原点を中心に正規化を行う。
3Dデータはパフォーマンスを向上させるため正規化させる必要があるとのこと。

スケーリングして頂点で学習するってことみたい。

・訓練データと検証データの作成

train_loader = DataLoader(ModelNet(modelnet_path, categories=categories,
   split='train', transform=transform, device=device),
   batch_size=12, shuffle=True)

val_loader = DataLoader(ModelNet(modelnet_path, categories=categories,
   split='test',transform=transform, device=device),
   batch_size=12)

"train_loader"訓練データ
"val_loader"検証データ

・分類器の訓練

ClassificationEngineクラスを使用して分類器の訓練を行う。 最も単純な点群ニューラルアーキテクチャであるPointNet分類器を用いて訓練および検証を行う。

engine = ClassificationEngine(PointNet(num_classes=len(categories)),
   train_loader, val_loader, device=device)
engine.fit()

ロジックを詳しく見よう思ったらドキュメントまだ書いてなかった。(2020/3/8時点)
とりあえず脳死で学習させる。

結果の可視化

検証データからデータをシャッフルし、新しいデータローダーを作成しサンプルバッチを取得する。

test_loader = DataLoader(ModelNet(modelnet_path, categories=categories,
                                split='test',transform=transform, device=device),
                       shuffle=True, batch_size=15)

モジュールのutils.pyをインポートし結果を可視化する。
/kaolin-master/examples/Classificationのフォルダに入っているのでパスを指定。

utils_path = '/content/drive/My Drive/xxxxx/kaolin-master/examples/Classification'
sys.path.append(utils_path)
from utils import visualize_batch

test_batch, labels = next(iter(test_loader))
preds = engine.model(test_batch)
pred_labels = torch.max(preds, axis=1)[1]
visualize_batch(test_batch, pred_labels, labels, categories)

GTが実際の3DモデルでPredが予測。今回可視化された15個についてはすべて正解で緑色になっている。
バッチリとイスとソファを分類する学習ができたと言える。

今回はソースコードに準じて学習を行った。混同行列とかF値とかは別の機会にやってみようかな。