ArcFace（ArcMarginProduct）


page:[AI開発]

リファレンス

ArcFaceの特徴

ArcFace Explained | Kaggleより
  • 似たものは近く、似てないものは遠くなるように学習する
  • 入力画像を右図のように写像する
    • 参加したコンペ では約90万件もあるので、左のような感じだと境目がはっきりしない

ArcFaceを使った距離学習モデル作成例

モデル全文

（75行）

class JPONet(nn.Module):

    def __init__(self,
                 n_classes,
                 model_name='efficientnet_b3',
                 use_fc=False,
                 fc_dim=224,
                 dropout=0.0,
                 loss_module='softmax',
                 s=30.0,
                 margin=0.50,
                 ls_eps=0.0,
                 theta_zero=0.785,
                 pretrained=True):
        """
        :param n_classes:
        :param model_name: name of model from pretrainedmodels
            e.g. resnet50, resnext101_32x4d, pnasnet5large
        :param pooling: One of ('SPoC', 'MAC', 'RMAC', 'GeM', 'Rpool', 'Flatten', 'CompactBilinearPooling')
        :param loss_module: One of ('arcface', 'cosface', 'softmax')
        """
        super(JPONet, self).__init__()
        print('Building Model Backbone for {} model'.format(model_name))

        self.backbone = timm.create_model(model_name, pretrained=pretrained)
        final_in_features = self.backbone.classifier.in_features

        self.backbone.classifier = nn.Identity() #恒等関数
        self.backbone.global_pool = nn.Identity()

        self.pooling = nn.AdaptiveAvgPool2d(1)

        self.use_fc = use_fc
        if use_fc:
            self.dropout = nn.Dropout(p=dropout)
            self.fc = nn.Linear(final_in_features, fc_dim)
            self.bn = nn.BatchNorm1d(fc_dim)
            self._init_params()
            final_in_features = fc_dim

        self.loss_module = loss_module
        if loss_module == 'arcface':
            self.final = ArcMarginProduct(final_in_features, n_classes, s=s, m=margin, easy_margin=False, ls_eps=ls_eps)
        elif loss_module == 'cosface':
            self.final = AddMarginProduct(final_in_features, n_classes, s=s, m=margin)
        elif loss_module == 'adacos':
            self.final = AdaCos(final_in_features, n_classes, m=margin, theta_zero=theta_zero)
        else:
            self.final = nn.Linear(final_in_features, n_classes)

    def _init_params(self):
        nn.init.xavier_normal_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0)
        nn.init.constant_(self.bn.weight, 1)
        nn.init.constant_(self.bn.bias, 0)

    def forward(self, x, label):
        feature = self.extract_feat(x)
        if self.loss_module in ('arcface', 'cosface', 'adacos'):
            logits = self.final(feature, label)
        else:
            logits = self.final(feature)
        return logits

    def extract_feat(self, x):
        batch_size = x.shape[0]
        x = self.backbone(x)
        x = self.pooling(x).view(batch_size, -1)

        if self.use_fc:
            x = self.dropout(x)
            x = self.fc(x)
            x = self.bn(x)

        return x
  • timmでEfficientNetを読み込み
  • ネットワークの終端を加工して、AcrFaceに渡す
    • 加工しているコードがとても参考になる

実体化（インスタンス作成）

（15行）

loss_module = 'arcface' #'cosface' #'adacos'
model_params = {
    'n_classes': df['cite_gid'].nunique(), #分類するクラス、コンペでは80万件に分類
    'model_name': 'efficientnet_b3', #利用モデル→b0、b3で悩む
    'use_fc': True, #GoogleColab Proではメモリ足らず、Trueで縮小
    'fc_dim': 224, #商標画像を224に前処理した
    'dropout': 0.0, #さわらず
    'loss_module': loss_module,
    's': 30.0, #さわらず
    'margin': 0.50, #さわらず
    'ls_eps': 0.0, #さわらず
    'theta_zero': 0.785, #さわらず
    'pretrained': True #初回のみ
    }
model = JPONet(**model_params)
  • 最終行がインスタンス作成しているところ
  • 「model_params」を引数にしてインスタンス作成
    • モデルのコンストラクタ（def __init__）にある引数＝パラメータをまとめて書いている
    • 「**」付けてる

作成されたモデルを確認してみる

見るところ

  • ネットワークの最後の部分
    • モデル実体の全文は長いので こちら
  • ベース（backbone）はEfficientNet（分類モデル）
  • ネットワークの最後が加工されている
  • ArcFace（ArcMarginProduct）が最後
    • （final）という名前になっている

ポイント

注目するところだけ抜き出した

(backbone): EfficientNet(・・・)
(pooling): AdaptiveAvgPool2d(output_size=1)
(dropout): Dropout(p=0.0, inplace=False)
(fc): Linear(in_features=1280, out_features=224, bias=True)
(bn): BatchNorm1d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(final): ArcMarginProduct()

  • 作成したインスタンスの中身（ネットワーク）を表示してから抜粋

ArcFaceをつなげる加工

  • 既存のネットワーク（ここではEfficientNet）にAcrFaceをつなげるため、既存のネットワークの一部を加工
  • AcrFace利用方法のチュートリアルが提供されていてそのまま使えるが、やり方を理解しておく
    • あぁ、こうやってネットワークつないでいくんだなってわかる

1つ目の加工

モデル追加箇所

        final_in_features = self.backbone.classifier.in_features
        self.backbone.classifier = nn.Identity() #恒等関数
        self.backbone.global_pool = nn.Identity()
  • モデル全文の28行目
  • (backbone)=EfficientNetの最後を加工するコード
    • 「classifier」「global_pool」を恒等関数に置き換える
    • 置き換える前に「in_features」を取り出しておく
      • 1280（＝EfficientNet-b0の圧縮後特徴量次元数）

加工前後（インスタンスの中身）

(global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))
(classifier): Linear(in_features=1280, out_features=1000, bias=True)

↓↓↓

(global_pool): Identity()
(classifier): Identity()

2つ目の加工

モデル追加箇所

        self.pooling = nn.AdaptiveAvgPool2d(1)
  • モデル全文の31行目
  • 上記が追加
  • (global_pool)の代わりと思われる
    • SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))をIdentity()に置き換えている
  • poolingの後にclassificationする形は、元のモデルから崩さないようにしているように見える

加工後追加されたもの（インスタンスの中身）

(pooling): AdaptiveAvgPool2d(output_size=1)

3つ目の加工

モデル追加箇所

        self.use_fc = use_fc
        if use_fc:
            self.dropout = nn.Dropout(p=dropout)
            self.fc = nn.Linear(final_in_features, fc_dim)
            self.bn = nn.BatchNorm1d(fc_dim)
            self._init_params()
            final_in_features = fc_dim
  • モデル全文の33行目
  • 次元削除
    • 使わない方が精度が当然上がる
    • GoogleColab Pro程度のメモリではオーバーフローしてしまうので、仕方なく利用
    • 既存ネットワークへのつなぎ方として勉強になるコード
  • 線形結合（nn.Linear）でfinal_in_features次元
    • fc_dim次元に変換
  • 6行目
    • モデル全文の51行目の「_init_params」を呼び出し
    • 追加したネットワーク分のパラメータを初期化している（後述）
  • 「self.use_fc」
    • 「forward」で使うのでインスタンス内変数（self）として保持

加工後追加されたもの（インスタンスの中身）

(dropout): Dropout(p=0.0, inplace=False)
(fc): Linear(in_features=1280, out_features=224, bias=True)
(bn): BatchNorm1d(224, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

追加モデルのパラメータ初期化箇所

    def _init_params(self):
        nn.init.xavier_normal_(self.fc.weight)
        nn.init.constant_(self.fc.bias, 0)
        nn.init.constant_(self.bn.weight, 1)
        nn.init.constant_(self.bn.bias, 0)
  • モデル全文の51行目
  • Linear（線形結合）
    • 「weight（重み）」「bias（バイアス）」を追加
    • バイアスは0
    • 重みはXavierの初期値
  • BatchNorm1（バッチ正規化）
    • 「weight（重み）」「bias（バイアス）」を追加
    • バイアスは0
    • 重み1

4つ目の加工

モデル追加箇所

        self.loss_module = loss_module
        if loss_module == 'arcface':
            self.final = ArcMarginProduct(final_in_features, n_classes, s=s, m=margin, easy_margin=False, ls_eps=ls_eps)
        elif loss_module == 'cosface':
            self.final = AddMarginProduct(final_in_features, n_classes, s=s, m=margin)
        elif loss_module == 'adacos':
            self.final = AdaCos(final_in_features, n_classes, m=margin, theta_zero=theta_zero)
        else:
            self.final = nn.Linear(final_in_features, n_classes)
  • モデル全文の41行目
  • ArcFace（ArcMarginProduct）登場
    • 「final」として「ArcMarginProduct」をつなげる
    • 引数は取得した「final_in_features」と分類したいクラスの数「n_classes」
  • Identifyにした部分がこれに置き換わったような感じ
    • 加工前は「(classifier): Linear(in_features=1280, out_features=1000, bias=True)」
  • 「use_fc」を使用している（True）場合
    • 「final_in_features」には「fc_dim」が設定されている

加工後追加されたもの（インスタンスの中身）

(final): ArcMarginProduct()

補足

損失関数について

以下は、損失計算にかかる部分のみコードを抜き出した

def fetch_loss():
    loss = nn.CrossEntropyLoss()
    return loss
criterion = fetch_loss()
model = JPONet(**model_params)
output = model(image,targets)
loss = criterion(output,targets)


AIまとめ

DOCKER/WSL

AI開発

ライブラリ

うんちく

Posted by futa