pytorchでの転移学習後の保存したモデルがロードができない。Grad-CAMの実装ができない。

Question

環境：Ubuntu 18.04
フレームワーク：pytorch

困っていること：
①作成したモデルのロードができない
②Grad-CAMを実装しようとすると以下のエラーが出る

初心者ながら頑張って、現在pytorchを使ってresnet18の転移学習を行っています。
実際に学習には成功し、推論、Grad-CAMの実装を試みています。
上記の2点で躓いています。ご教授いただけたら幸いです。

①
様々なサイトに乗っている以下のコードを実行

torch.save(the_model, PATH)

the_model = torch.load(PATH)

上記のような保存、ロードは問題なくできました。
ですが、下記で保存したモデルをうまくロードできません。

torch.save(the_model.state_dict(), PATH)

さまざまなサイトで下記が記載してありますが、name 'TheModelClass' is not defined
とエラーが出ます。

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

*args、*kwargsについても調べて理解はしたつもりですが、TheModelClassをどう定義するのか？などわかりません。
ご教授お願いします。

②Grad-CAMについて

サイトに掲載されたソースコードを転用して、自分の作成したモデルで試そうとするのですが、
Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the sameというエラーが出ます。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from torchvision import datasets, transforms
from torchvision import models
from tqdm import tqdm_notebook as tqdm
from PIL import Image
import cv2


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class GradCAM:
    def __init__(self, model, feature_layer):
        self.model = model
        self.feature_layer = feature_layer
        self.model.eval()
        self.feature_grad = None
        self.feature_map = None
        self.hooks = []

        # 最終層逆伝播時の勾配を記録する
        def save_feature_grad(module, in_grad, out_grad):
            self.feature_grad = out_grad[0]
        self.hooks.append(self.feature_layer.register_backward_hook(save_feature_grad))

        # 最終層の出力 Feature Map を記録する
        def save_feature_map(module, inp, outp):
            self.feature_map = outp[0]
        self.hooks.append(self.feature_layer.register_forward_hook(save_feature_map))

    def forward(self, x):
        return self.model(x)

    def backward_on_target(self, output, target):
        self.model.zero_grad()
        one_hot_output = torch.zeros([1, output.size()[-1]])
        one_hot_output[0][target] = 1
        output.backward(gradient=one_hot_output, retain_graph=True)

    def clear_hook(self):
        for hook in self.hooks:
            hook.remove()

model.full=torch.load("PATH")

model.load_state_dict(torch.load(model_path))
image_model.eval()
id_to_label = {
    0: 'other',
    1: 'true'
}

grad_cam = GradCAM(model=image_model, feature_layer=list(image_model.layer4.modules())[-1])


from PIL import Image
from torchvision.transforms.functional import to_pil_image

VISUALIZE_SIZE = (224, 224)  
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

image_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

path = "PATH2"
image = Image.open(path)
image.thumbnail(VISUALIZE_SIZE, Image.ANTIALIAS)
display(image)

# save image origin size
image_orig_size = image.size # (W, H)

img_tensor = image_transform(image)
img_tensor = img_tensor.unsqueeze(0)

model_output = grad_cam.forward(img_tensor)
target = model_output.argmax(1).item()

上記実行した際ののエラー

Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

よろしくお願いします。

Answer 1

①は、
Pytorchでモデルの保存と読み込み
の「保存時はstate_dict()を使う」を見てください

Comments

[tomotomotomo777]

ご回答ありがとうございます。
助言いただいたサイトを参考にして以下の通りに保存し、何とかロードできました。

modelsメソッドで雛形となるモデルを読みだして、
保存した重みを当てはめるという認識なんですが、
私の認識は正しいでしょうか。

```python
from torchvision import models
model = models.resnet18()
torch.save(model.to('cpu').state_dict(), 'model.pth')
```

```python
from torchvision import models
model = modelsresnet18()
model.load_state_dict(torch.load('model.pth'))
```