ご存知のとおり、大量のデータで事前トレーニングされたバックボーンネットワークを使用すると、モデル全体の汎化能力が向上しますが、ネットワークのバックボーンネットワークを置き換える場合、元の重みを直接使用することはできません。このプロジェクトの目的は、バックボーンネットワークを置き換えた後にネットワークの事前トレーニングの重みを「盗む」ことです。

結論は次のとおりです。DeeplabV3+ のバックボーンネットワークを Xception によって mobilenetv3 に置き換え、事前トレーニングと適用外の事前トレーニングを使用します。最初のエポックと最初の 171 エポックの結果、最初の 200 ラウンドの損失関数は次のとおりです。次のように：

不适用预训练：
--->1轮miou如下:
   mIoU: 16.11; mPA: 16.67; Accuracy: 96.68
使用预训练：
--->1轮miou如下:
   mIoU: 40.53; mPA: 56.8; Accuracy: 97.09

不适用预训练：
--->171轮miou如下:
   mIoU:64.68 ; mPA: 78.0; Accuracy: 98.82
使用预训练：
--->171轮miou如下:
   mIoU:86.08 ; mPA: 92.54; Accuracy: 99.56

事前に訓練された重みを変換してロードした後、初期の勾配降下と最終モデルのパフォーマンスの両方が大幅に改善された ことがわかります。

1. ウェイトファイルの構造

pyotrch の重みファイルは .pth です。これは、シリアル化されたコレクションであるcollections.OrderedDictで構成されます。

これは、レイヤーの名前 + レイヤーの Tensor ウェイトの 2 つの要素からなるユニットで構成されます。これを横切ることで、2 つを別々に取り出すことができます。

for k, v in pretrained_dict.items():

検証の結果、kのデータ型はstr型です。元のアダプティブローディングコードでは、他のネットワークの事前学習済み重みが使えない理由は、重み名が一致しません。重み名がstr型のみであるためです。、replace を使用すると、矛盾した部分が変更された場合、ウェイトファイルをネットワークに追加できます。

2. 重み名を変更する

まず、この方法は重み名に何らかの関係がある場合のみですので、バックボーンネットワークを移植する場合は、関数名や構造名が元のネットワークと一致するようにしてください。バックボーンネットワークはモデルの先頭にあるので、最初の層を出力することで、変更後のネットワークと元のネットワークの違いが分かります。この例では、mobilenetv3 と bubbliiiiing の YoloV4 の deeplabV3 + を素材として使用しています。 mobilenentv3 ネットワークを deeplabv3+ に移動します。

OrderedDictのレイヤーをpython3で直接出力すると、シリアル化されたデータ構造なのでエラーが報告されますので、list(.)を使って非シーケンシャルなデータ構造に変換して出力する必要があります。出力を完了します。

print('源模型文件格式为:{}'.format(list(pretrained_dict.keys())[0]))
print('目标模型文件格式为:{}'.format(list(model_dict.keys())[0]))

出力後、2 つのモデル間のギャップは次のようになります。

2 つの変数を使用して 2 つの変数の差を一時的に保存し、これら 2 つの変数を replace() 関数に送信してレイヤー名の変更を完了します。上図に示すように、ソースモデルとターゲットモデルの間には、ソースモデルのモデルがターゲットモデルよりも1 つ多いというギャップがあるため、それを削除するだけです。

3. 変換して保存する

この例で使用されている関数は個別にパッケージ化されています。

def model_converter(custom_model ,model_path):
    model_dict      = custom_model.state_dict()
    pretrained_dict = torch.load(model_path, map_location = torch.device('cpu'))
    load_key, no_load_key, temp_dict = [], [], {}
    #  展示骨干网络的第一层
    print('源模型文件格式为:{}'.format(list(pretrained_dict.keys())[0]))
    print('目标模型文件格式为:{}'.format(list(model_dict.keys())[0]))
    print('请将两模型之间不同的部分输入:')
    orgStr = input('源模型:')
    targetStr = input('目标模型:')

    print('--->开始模型转换')
    for k, v in pretrained_dict.items():
        k = k.replace(orgStr,targetStr)
        if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
            temp_dict[k] = v
            load_key.append(k)
        else:
            no_load_key.append(k)
    #  将权重更新到模型中
    model_dict.update(temp_dict)
    custom_model.load_state_dict(model_dict)
    #  保存模型
    torch.save(custom_model.state_dict(), 'converted_weights.pth')

この関数にインスタンス化したモデルとソースモデルファイルを入力するだけで、修正後、一致した部分は自動変換され、一致しない部分はno_load_key に一時的に保存され、単体で解析することができます。テスト後、この方法は基本的にバックボーンネットワークの変換を完了できます。

[PyTorch] 事前トレーニングされた重みの変換

1. ウェイトファイルの構造

2. 重み名を変更する

3. 変換して保存する

おすすめ