SegViT 公式モデルのソース コードから学習して、独自のローカル コード ファイルにデプロイしたい
1. 環境構築
公式 Web サイトでは、mmcv-full=1.4.4 および mmsegmentation=0.24.0 のインストールが必要です。
そうする前に、元のバージョンの mmcv および mmsegmentation をアンインストールしてください。
pip uninstall mmcv
pip uninstall mmcv-full
pip uninstall mmsegmentation
mmcv をインストールする
このうち mmcv には、mmcv の完全版 (当初は mmcv-full と呼ばれていました) と、mmcv-lite の合理化版 (当初は mmcv と呼ばれていました) の 2 つのバージョンがあり、バージョン 2.0.0 以降は名前が変更されました。具体的な違いについては、 mmcv 公式 Web サイトのマニュアルを参照してください。また、mmcv-full (つまり、mmcv の完全版) をインストールするブログは、主に mmcv 公式 Web サイトのマニュアルを参照してください。mmcv>=2.0.0 をインストールしたい場合は、詳細を説明せずに、公式 Web サイトのマニュアルに従って直接インストールできます。たとえば、私が mmcv-full==1.4.4 をインストールしたなど、過去のバージョンをインストールしたい場合は、私の記録を参照してください。mmcv をインストールする前に、まず pytorch と cuda の対応するバージョンを知っておく必要があります。pytorch のバージョンを表示します。
python -c 'import torch;print(torch.__version__)'
バージョン情報が出力されていれば、pytorchがインストールされていますので、cuda
のバージョンを確認してください：お使いの環境
のpytorchに対応したcudaのバージョンを確認してください。これは、pytorch に対応する cuda のバージョンを確認するために使用したコマンドです。
python -c 'import torch;print(torch.version.cuda)'
次のように書くこともできます。
参考ブログ：https://blog.csdn.net/qq_49821869/article/details/127700187
python
>>>import torch
>>>torch.version.cuda
ここで、pytorch のバージョンは 1.11.0 である必要があり、対応する cuda のバージョンは 11.3 です。
参考ブログ：https://blog.csdn.net/qq_41661809/article/details/125345690
そこで、次のコマンドを入力しました。
pip install mmcv-full==1.4.4 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.11.0/index.html
失敗したので、この Web サイトにアクセスして確認したところ、使用できる最低バージョンは 1.4.7 であることがわかったので、
コマンドを次のように変更しました。
pip install mmcv-full==1.4.7 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.11.0/index.html
mmcv-full のインストールが完了しました
mmセグメンテーションをインストールする
最初は公式 Web サイトの手順に従って mmsegmentation をインストールしました
が、mmcv>=2.0.0 が必要で、インストールされたバージョンは mmsegmentation==1.0.0 であり、要件と矛盾していました。
mmsegmentation は mmcv バージョンと一致する必要があることに注意してください。
参考ブログ：https://blog.csdn.net/CharilePuth/article/details/122909620
そこで私は直接こう言いました。
pip install mmsegmentation==0.24.0
インストール成功。
「Pip のインストール パッケージは、水を飲むのと同じくらい簡単です」 - かつて偉い人は言いました。
2. コード！
モデル構成ファイルの検索
公式 Web サイトにアクセスし、「トレーニング:ハイライト」でモデルに対応する構成ファイルを見つけます。この記事のハイライトの 1 つは計算コストを削減できる収縮構造であることがわかり
ました。そこで、次に収縮構造を選択します。
実行したいイメージのサイズは 512 であるため、このコードの結果で、同じ 512*512 を持つ COCO データ セットの対応するモデルが見つかりました。
configs フォルダーに戻って、このデータ セットに対応するネットワーク モデルを見つけます。
コードを観察して、使用されているバックボーンが vit_shrink で、デコード ヘッドが TPNATMHead であることを確認します。
パラメーター設定に注意し、構成ファイルにも注意してください。 __base__ のパラメータ。モデルを入力すると、そのパラメータがモデル内で宣言されます。
モデルコードを探す
バックボーン フォルダーに入り、vit_shrink ネットワークを見つけて、
それをコピーして自分の py ファイルに貼り付け
、decode_heads フォルダーでデコード ヘッダー コードを見つけて、
それをコピーして自分の py ファイルに貼り付けます。
コードをパッチアップする
- ライブラリ
ファイルに不足しているものをライブラリ ファイルに追加します。たとえば、tpn_atm_head デコーダ コード内の他の 2 つのデコーダ コードの内容を参照する必要がある場合は、Ctrl C+V キーを押しながら他の 2 つのデコーダ コードを直接入力します。使用する必要があるモジュールはそのままにしておきます。
- 入出力バックボーンの入出力を確認する
:
デコーダ部分の入出力は図のようになります:
SegViT を書いて入出力をテストします 設定ファイルを参照して、対応する設定を事前に宣言してください:
class SegViT(nn.Module):
def __init__(self, num_class):
super(SegViT, self).__init__()
out_indices = [7,23]
in_channels = 1024
img_size = 512
# checkpoint = './pretrained/vit_large_p16_384_20220308-d4efb41d.pth'
checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_large_p16_384_20220308-d4efb41d.pth'
# self.backbone = get_vit_shrink()
self.backbone = vit_shrink(
img_size=(img_size,img_size),embed_dims=in_channels,num_layers=24,drop_path_rate=0.3,num_heads=16,out_indices=out_indices)
self.decoder = TPNATMHead(
img_size=img_size,in_channels=in_channels,channels=in_channels,embed_dims=in_channels//2,num_heads=16,num_classes=num_class,num_layers=3, use_stages=len(out_indices))
def forward(self, _x):
x = self.backbone(_x)
out = self.decoder(x)
# if self.training:
# return out['pred'], out['ce_aux']
# else:
# return out
return out
チェックアウトの種類を実行する
if __name__ == "__main__":
x = torch.randn(4, 3, 512, 512)
net = SegViT(6)
# flops, params = profile(net, (x,))
# print('flops: %.2f G, params: %.2f M' % (flops / 1000000000.0, params / 1000000.0))
# res, aux = net(x)
res = net(x)
print(res)
すると、出力は辞書型、予測はキー名 pred に対応する値、値はテンソル型、形状サイズは (4,6,512,512) で、出力が正しいことがわかります。
次に、補助ブランチの出力を見つける必要があります。
デコーダ ヘッダーの前方にあります:
コメントを削除して補助ブランチの出力を取得します (補助ブランチの出力は辞書要素の形式で atm_out に追加され、デバッグできます)。コメントを追加することを忘れないでください。このうち
、1枚のカードで実行しているため、SyncBNをBNに変更しました。そうしないとエラーが報告されます。
さらに、トレーニング フェーズとテスト フェーズの出力は異なり、デバッグして確認できます。
def forward(self, _x):
x = self.backbone(_x)
out = self.decoder(x)
if self.training:
return out['pred'], out['ce_aux']
else:
return out
- 重量ファイルの読み込み なお
、重量ファイルは事前にダウンロードしておくことも可能です。
def get_vit_shrink(pretrained=True, img_size=512, in_channels=1024, out_indices=[7,23]):
model = vit_shrink(
img_size=(img_size,img_size),embed_dims=in_channels,num_layers=24,drop_path_rate=0.3,num_heads=16,out_indices=out_indices)
if pretrained:
checkpoint = '权重文件所在路径'
# if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict']
# else: state_dict = checkpoint
model.load_state_dict(checkpoint, strict=False)
return model
最終的なモデルは次のとおりです。
class SegViT(nn.Module):
def __init__(self, num_class):
super(SegViT, self).__init__()
out_indices = [7,23]
in_channels = 1024
img_size = 512
# checkpoint = './pretrained/vit_large_p16_384_20220308-d4efb41d.pth'
# checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_large_p16_384_20220308-d4efb41d.pth'
self.backbone = get_vit_shrink()
self.decoder = TPNATMHead(
img_size=img_size,in_channels=in_channels,channels=in_channels,embed_dims=in_channels//2,num_heads=16,num_classes=num_class,num_layers=3, use_stages=len(out_indices))
def forward(self, _x):
x = self.backbone(_x)
out = self.decoder(x)
if self.training:
return out['pred'], out['ce_aux']
else:
return out
- 最終的な入出力を確認して
終了です。
3. モデルを実行する
独自のフレームワークでパラメータを設定して実行します。
仕上げる。