普段Tensorflow(Keras)を使って機械学習をしていますが、後方互換性の無さに嫌気が差したのでPyTorchを使ってみました。
今回はPyTorchに組み込まれているMask R-CNNを自前のデータセットで学習させる上で参考になる資料が少なかったため、主にデータセットを構築していくところに焦点を当てています。

#動作環境と用意するもの

Azure Machine Learning(データセットの登録・呼び出しや実験の構築等、特有の書き方は出てきません。ほぼローカルで動くのと同じPython Script上での話になります)
Segmentationしたい画像とそのアノテーションデータ(今回はCOCO形式のデータを使います。LabelMe等で自作していることを前提とします)

#参考にした記事

TORCHVISION OBJECT DETECTION FINETUNING TUTORIAL
- 結局ここに書かれていることが全てです。翻訳したものも検索結果には出てきますが、訳が曖昧なのでこちらを読みましょう。
- とはいってもここで書かれていることだけではCOCO形式のデータを読み込ませることはできません。
【PyTorchチュートリアル⑧】TorchVision Object Detection Finetuning Tutorial
- この記事の中ほどにあるデータセットの出力形式が参考になりました。

#COCO形式のデータセットをPyTorchに読み込ませる方法
上記のチュートリアルでは画像とマスク画像のセットをデータセットとして読み込ませる方法が記載されています。
そのため、COCO形式のデータセットを読み込ませるためにはここに書かれていること以外の工夫が必要になります。

工夫をしていく前に今回読み込ませるデータセットを紐解いていきましょう。

まず、COCO形式でのアノテーションデータはJsonファイルに格納されています。(中身の詳細はこちらを参照)
PyTorch(TorchVision)の公式ドキュメントを読むとTorchVision.DatasetにCOCOと書かれているクラスが存在する(COCODetecitonとCOCOCaptionが書かれている)ことから、これを使えばDataLoaderにデータを渡して学習させることが可能なように思えます。

ですがこれを使っても学習させることは出来ません。
というのも画像データ・アノテーションデータのどちらもTensorに変換しなければいけないのが公式チュートリアルのコードから読み取れますがこの処理をCOCODetecitonはしてくれません。(当該コード)
その変換処理はself.transformsに渡す関数が行ってくれます。
画像(PIL形式に変換したもの)はこの関数がTensorに変換してくれますが、target(当該画像に対応するアノテーションマスクの情報が入った辞書を要素とするリスト)をTensorに変換する関数はPyTorch側が用意していません。ですので自作する必要があります。(じゃあなぜCOCODetection等が存在するのかは不明)

自作したコード

コンストラクタ等、カスタマイズする必要のないものは継承してきたCocoDetectionのものを使います。(CocoDetectionを使うためにはPycocotoolsを別途pipやcondaを使ってインストールする必要があります)

class CustomDataset(CocoDetection):
    def __getitem__(self, index):
        #継承した親メソッドから必要な変数を受け取る
        idx = self.ids[index]
        image = self._load_image(idx)
        target = self._load_target(idx)

        #target(アノテーションデータ)をデータセットとして読み込めるように変換する      
        boxes = [x['bbox'] for x in target]
        labels = [x['category_id'] for x in target]
        image_id = idx
        area = [box[2] * box[3] for box in boxes]
        iscrowd = [x['iscrowd'] for x in target]
        masks = [self.coco.annToMask(x) for x in target]

        targets = {}
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        targets["boxes"] = torchvision.ops.box_convert(boxes,'xywh','xyxy')
        targets["labels"] = torch.as_tensor(labels, dtype=torch.int64)
        targets["masks"] = torch.as_tensor(masks, dtype=torch.uint8)
        targets["image_id"] = torch.tensor([image_id])
        targets["area"] = torch.tensor(area)
        targets["iscrowd"] = torch.as_tensor(iscrowd, dtype=torch.int64)

        #transformsを受け取っていれば内容に従って変換する
        if self.transforms is not None:
            image, targets = self.transforms(image, targets)

        return image, targets

コードの大体は公式チュートリアルに書かれているものがベースです。
先述の通り、targetは辞書が要素のリストなのでリスト内包記法を使って各要素を抽出したあとTensorに変換します。

ここにも罠がひとつあって、コードを読むとわかるのですが、COCO形式のkeyとデータセットに読み込ませる際のkeyがイコールではありません。(bbox→boxes等)
なのでそのままだと読み込めません(なんでその処理をPyTorch側でやってくれないのかがわかりません)

key名を変換するだけの要素もありますが、areaはbboxから面積を計算する必要がある等の一工夫が必要です。

またbbox(bounding box)の形式ですが、COCO形式だと[xmin,ymin,width,height]なのに対して、PyTorch(R-CNN)側は[xmin,ymin,xmax,ymax]であることが前提ですが、この変換処理をしてくれる関数は存在します。(なんでこれだけあるんだ)

加えてmasks(segmentation mask)も形式が異なるので変換が必要です。
COCO形式ではポリゴン(x,yの点群情報)でmaskを形成しているのに対して、PyTorchではMask画像(0~1に正規化した画像情報)を想定していますので、この変換も必要です。
なお、この処理はPycocotoolsがやってくれますので、難しい処理を自作する必要はなさそうです。
この後、このクラスを呼び出してDataLoaderに読み込ませれば学習を実行してくれます。(学習コードは公式チュートリアルと同一です)

#終わりに
PyTorchを使うのは初めてでしたが、全体的な使いやすさはTensorflowよりも上だなと思いました。
ただネットワークを実装したり、層を自分で追加したりという真髄の部分は今回書いていないので、PyTorchの本領を発揮とまではいきません。

一方、今回やろうとしたことの情報の少なさのせいか、動かせるようになるまで1週間ほど掛かりました。まだまだ精進が必要です。