Google Colaboratory 上で重みモデル.pthをロードできない

Question

実現したいこと

Google colaboratory上で自前の重みモデルkanji.pthをロードし学習したいです。
くずし字の漢字を検出・認識するプログラムを作っています。
現在
Create Pytorch Dataset for Classifying Characters
内にあるDemoModel↓

class DemoModel(torch.nn.Module):
    def __init__(self):
        super(DemoModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=1,
                                     out_channels=16,
                                     kernel_size=7)
        self.relu1 = torch.nn.ReLU(inplace=True)
        self.maxpool1 = torch.nn.MaxPool2d(kernel_size=2)
        self.conv2 = torch.nn.Conv2d(in_channels=16,
                                     out_channels=128,
                                     kernel_size=6)
        self.relu2 = torch.nn.ReLU(inplace=True)
        self.maxpool2 = torch.nn.MaxPool2d(kernel_size=2)
        self.fc = torch.nn.Linear(in_features=128*8*8,
                                  out_features=4212,
                                  bias=True)
        self.log_softmax = torch.nn.LogSoftmax(dim=-1)

    def forward(self, x):
        out = self.conv1(x) # (batch, 1, 48, 48) -> (batch, 16, 42, 42)
        out = self.relu1(out)
        out = self.maxpool1(out) # (batch, 16, 42, 42) -> (batch, 16, 21, 21)
        out = self.conv2(out) # (batch, 16, 21, 21) -> (batch, 128, 16, 16)
        out = self.relu2(out)
        out = self.maxpool2(out) # (batch, 128, 16, 16) -> (batch, 128, 8, 8)
        out = out.view(out.size(0), -1) # (batch, 128, 8, 8) -> (batch, 8192)
        out = self.fc(out) # (batch, 8192) -> (batch, 4212)
        out = self.log_softmax(out)
        return out

を
Create Pytorch Dataset for Classifying Charactersに書かれている方法で学習した後、
torch.save(DemoModel().state_dict(),'/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/' + 'kanji' + '.pth')
で保存したモデルで物体検出を行おうとしたところ以下のエラーメッセージが発生しました。

発生している問題・エラーメッセージ

Loading weights into state dict...
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-19-88fdfbbc72bc> in <module>()
     17 net = build_ssd('test', 300, 4213)
     18 #net = build_ssd('test', 300, 21)
---> 19 net.load_weights('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/kanji.pth')
     20 #net.load_weights('/content/drive/MyDrive/Colab Notebooks/chapter7/weights/ssd300_mAP_77.43_v2.pth')
     21 

1 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1405         if len(error_msgs) > 0:
   1406             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1407                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1408         return _IncompatibleKeys(missing_keys, unexpected_keys)
   1409 

RuntimeError: Error(s) in loading state_dict for SSD:
	Missing key(s) in state_dict: "vgg.0.weight", "vgg.0.bias", "vgg.2.weight", "vgg.2.bias", "vgg.5.weight", "vgg.5.bias", "vgg.7.weight", "vgg.7.bias", "vgg.10.weight", "vgg.10.bias", "vgg.12.weight", "vgg.12.bias", "vgg.14.weight", "vgg.14.bias", "vgg.17.weight", "vgg.17.bias", "vgg.19.weight", "vgg.19.bias", "vgg.21.weight", "vgg.21.bias", "vgg.24.weight", "vgg.24.bias", "vgg.26.weight", "vgg.26.bias", "vgg.28.weight", "vgg.28.bias", "vgg.31.weight", "vgg.31.bias", "vgg.33.weight", "vgg.33.bias", "L2Norm.weight", "extras.0.weight", "extras.0.bias", "extras.1.weight", "extras.1.bias", "extras.2.weight", "extras.2.bias", "extras.3.weight", "extras.3.bias", "extras.4.weight", "extras.4.bias", "extras.5.weight", "extras.5.bias", "extras.6.weight", "extras.6.bias", "extras.7.weight", "extras.7.bias", "loc.0.weight", "loc.0.bias", "loc.1.weight", "loc.1.bias", "loc.2.weight", "loc.2.bias", "loc.3.weight", "loc.3.bias", "loc.4.weight", "loc.4.bias", "loc.5.weight", "loc.5.bias", "conf.0.weight", "conf.0.bias", "conf.1.weight", "conf.1.bias", "conf.2.weight", "conf.2.bias", "conf.3.weight", "conf.3.bias", "conf.4.weight", "conf.4.bias", "conf.5.weight", "conf.5.bias". 
	Unexpected key(s) in state_dict: "conv1.weight", "conv1.bias", "conv2.weight", "conv2.bias", "fc.weight", "fc.bias". 

該当のソースコード

PyTorch 学習済みモデルでサクッと物体検出をしてみる
上記のサイトを参考にさせていただきました。
下記は学習済みモデルkanji.pthを使用して物体検出をするプログラムです。
build_ssdコマンドに与えている4213という数字は今回検出する漢字の種類+1の値です。(4212+1)
漢字の種類の個数については
Create Pytorch Dataset for Classifying Characters内にあるuc_list(漢字を検出するためのUNICODEのリストです)の個数をpandasを使用し確認しました。

import os
import sys
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)
 
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import cv2
import glob
from ssd import build_ssd
from matplotlib import pyplot as plt
 
# SSDモデルを読み込み
net = build_ssd('test', 300, 4213)   
#net = build_ssd('test', 300, 21)   
net.load_weights('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/kanji.pth')
#net.load_weights('/content/drive/MyDrive/Colab Notebooks/chapter7/weights/ssd300_mAP_77.43_v2.pth')
 
# 関数 detect    
def detect(image, count):
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    x = cv2.resize(image, (300, 300)).astype(np.float32)  # 300*300にリサイズ
    x -= (104.0, 117.0, 123.0)
    x = x.astype(np.float32)
    x = x[:, :, ::-1].copy()
    x = torch.from_numpy(x).permute(2, 0, 1)  # [300,300,3]→[3,300,300]
    xx = Variable(x.unsqueeze(0))     # [3,300,300]→[1,3,300,300]    
    # 順伝播を実行し、推論結果を出力    
    y = net(xx)
 
    from data import VOC_CLASSES as labels
    plt.figure(figsize=(10,6))
    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    plt.imshow(rgb_image)
    currentAxis = plt.gca()
    # 推論結果をdetectionsに格納
    detections = y.data
    # scale each detection back up to the image
    scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)
    
    # バウンディングボックスとクラス名を表示
    for i in range(detections.size(1)):
        j = 0
        # 確信度confが0.6以上のボックスを表示
        # jは確信度上位200件のボックスのインデックス
        # detections[0,i,j]は[conf,xmin,ymin,xmax,ymax]の形状
        while detections[0,i,j,0] >= 0.6:
            score = detections[0,i,j,0]
            label_name = labels[i-1]
            display_txt = '%s: %.2f'%(label_name, score)
            pt = (detections[0,i,j,1:]*scale).cpu().numpy()
            coords = (pt[0], pt[1]), pt[2]-pt[0]+1, pt[3]-pt[1]+1
            color = colors[i]
            currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
            currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor':color, 'alpha':0.5})
            j+=1
    plt.savefig('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/detect_img/'+'{0:04d}'.format(count)+'.jpg')
    plt.close()
 
def main():
    files = sorted(glob.glob('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/image_dir/*.jpg'))
    count = 1
    for i, file in enumerate (files):
        image = cv2.imread(file, cv2.IMREAD_COLOR)          
        detect(image, count)
        print(count)
        count +=1
 
if __name__ == '__main__':
    main()  

試したこと

net = build_ssd('test', 300, 4213)
を削除し
net = DemoModel()
を実行しロードを試みたのですが、

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-20-b619b3184e9c> in <module>()
     18 #net = build_ssd('test', 300, 4213)
     19 #net = build_ssd('test', 300, 21)
---> 20 net.load_weights('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/kanji.pth')
     21 #net.load_weights('/content/drive/MyDrive/Colab Notebooks/chapter7/weights/ssd300_mAP_77.43_v2.pth')
     22 

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in __getattr__(self, name)
   1129                 return modules[name]
   1130         raise AttributeError("'{}' object has no attribute '{}'".format(
-> 1131             type(self).__name__, name))
   1132 
   1133     def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None:

AttributeError: 'DemoModel' object has no attribute 'load_weights'

load_weights属性をDemoModelは持ってないと出てしまいました。

動作環境

Google Colaboratory
Python 3.7.11

Comments

[jbpb0]

> net = DemoModel()
を実行しロードを試みたのですが、

> load_weights属性をDemoModelは持ってないと出てしまいました。

net = DemoModel()
net.load_state_dict(torch.load('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/kanji.pth'))

ではダメでしょうか？

参考
https://tzmi.hatenablog.com/entry/2020/03/05/222813

[sugar_s]

jbpb0様返信ありがとうございます。
net.load_weightsをnet.load_state_dict(torch.load(...に変えたところパラメータのロードが通りました！
ありがとうございます。load_weightsに絞って検索をかけていたのでjbpb0様の知見のおかげで助かりました。
ただ、
RuntimeError: Given groups=1, weight of size [16, 1, 7, 7], expected input[1, 3, 300, 300] to have 1 channels, but got 3 channels instead
と新たな問題も出てしまったので、RGB画像をグレースケール画像にしてから検出できるよう頑張ります！

Answer 1

net = DemoModel()
net.load_state_dict(torch.load('/content/drive/MyDrive/Colab Notebooks/Kuzushiji_Visualisation/kanji.pth'))