RuntimeError: ResNet の state_dict のロード中にエラーが発生しました: state_dict にキーがありません: “conv1.0...

プロジェクトのシナリオ:

マルチ GPU 環境で Pytorch を使用してトレーニングされた Resnet 分類ネットワーク

問題の説明

畳み込みニューラルネットワーク ResNet のトレーニング後、テスト環境またはテストコードが単一の GPU バージョンまたは CPU バージョンを使用している場合、ネットワークのロード時にエラーが報告されます。エラーコードは次のとおりです。

net.load_state_dict(torch.load(args.weights))

エラーは次のように報告されます。

RuntimeError: Error(s) in loading state_dict for ResNet: 
	Missing key(s) in state_dict: "conv1.0.weights", "conv1.1.weights", "conv1.1.bias", ...

原因分析:

このエラーの主な理由は、state_dict がモデルの重みを読み込むときにパラメーターが一致しないことです。PyTorch のバージョン環境が一致していない、torch.nn.DataParallel() キーワードが一致していない、トレーニング環境とテスト環境の GPU が異なるなどが考えられます。

この種のエラーが発生しました。1 回目は GPU がトレーニング中、CPU がテスト中だったためです。もう 1 回目は、複数の GPU でトレーニングしており、トレーニング中のマルチ GPU 処理ではなく、テスト中に GPU 部分が処理されました。シングルGPUでした。

解決：

情况一: マルチ GPU トレーニング、テスト中はコードが単一 GPU である
解決策: 環境内で GPU を複数の GPU として定義し、キーワード torch.nn.DataParallel() を定義します。

import torch.nn as nn
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"

if __name__ == '__main__':
	# 此处为测试代码对接收参数的初始化
	parser = argparse.ArgumentParser()
	parser.add_argument('-gpu', action='store_true', default=True, help='use gpu or not')
	......
	......
	......

	# 加载网络
	net = get_network(args)
	
	# 以下为多GPU环境增加代码，以解决本博客中提到的报错问题
	net = net.cuda()
	net = nn.DataParallel(net, device_ids=[0, 1, 2, 3])
	
	# 以下为原始测试代码部分
	net.load_dict(torch.load(args.weights))
	net.eval()
	......
	......
	......

情况二:GPU トレーニング、テスト時はコードが CPU
解決策: モデルの重みをロードするときに、map_location='cpu'

if __name__ == '__main__':
	# 此处为测试代码对接收参数的初始化
	parser = argparse.ArgumentParser()
	parser.add_argument('-gpu', action='store_true', default=True, help='use gpu or not')
	......
	......
	......

	# 加载网络
	net = get_network(args)
	
	
	# 以下为原始测试代码部分
	# net.load_dict(torch.load(args.weights))
	# 改为cpu模式
	net.load_dict(torch.load(args.weights, map_location='cpu'))
	net.eval()
	......
	......
	......

ケース 2 で発生する可能性のあるエラーメッセージは次のとおりです。

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

参考ブログ投稿:

torch.load を map_location=torch.device('cpu') とともに使用して、ストレージを CPU にマッピングしてください。

トレーニングパラメータとテストパラメータの間に不一致がある場合は、次のブログ投稿を参照してください。

【错误记录】RuntimeError: DataParallel の state_dict のロード中にエラーが発生しました: モジュールのサイズが一致しません