【Pytorch】配布されているモデルの読み込み方が分からない

Question

実現したいこと

Pytorch初心者です。
以下のモデルを使って画像のスコアを出力したいのですが、エラーが出てしまいます。

使用したいモデル
https://huggingface.co/Eugeoter/waifu-scorer-v3

発生している問題・分からないこと

・発生しているエラー
RuntimeError: Error(s) in loading state_dict for MLP: Missing key(s) in state_dict: "layers.7.weight", "layers.7.bias". Unexpected key(s) in state_dict: "layers.8.weight", "layers.8.bias", "layers.10.weight", "layers.10.bias", "layers.10.running_mean", "layers.10.running_var", "layers.10.num_batches_tracked", "layers.12.weight", "layers.12.bias", "layers.14.weight", "layers.14.bias", "layers.14.running_mean", "layers.14.running_var", "layers.14.num_batches_tracked", "layers.16.weight", "layers.16.bias", "layers.18.weight", "layers.18.bias", "layers.2.running_mean", "layers.2.running_var", "layers.2.num_batches_tracked", "layers.6.running_mean", "layers.6.running_var", "layers.6.num_batches_tracked". size mismatch for layers.0.weight: copying a param with shape torch.Size([2048, 768]) from checkpoint, the shape in current model is torch.Size([1024, 768]). size mismatch for layers.0.bias: copying a param with shape torch.Size([2048]) from checkpoint, the shape in current model is torch.Size([1024]). size mismatch for layers.2.weight: copying a param with shape torch.Size([2048]) from checkpoint, the shape in current model is torch.Size([128, 1024]). size mismatch for layers.2.bias: copying a param with shape torch.Size([2048]) from checkpoint, the shape in current model is torch.Size([128]). size mismatch for layers.4.weight: copying a param with shape torch.Size([512, 2048]) from checkpoint, the shape in current model is torch.Size([64, 128]). size mismatch for layers.4.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([64]). size mismatch for layers.6.weight: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([16, 64]). size mismatch for layers.6.bias: copying a param with shape torch.Size([512]) from checkpoint, the shape in current model is torch.Size([16]).

こちらのソースコードを使用しました（モデルの説明文に記載されていたため）
https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/simple_inference.py

該当のソースコード

import numpy as np
import torch
import pytorch_lightning as pl
import torch.nn as nn
from torchvision import datasets, transforms
import clip
from PIL import Image, ImageFile


#####  This script will predict the aesthetic score for this image file:

img_path = "anime.jpg"

# if you changed the MLP architecture during training, change it also here:
class MLP(pl.LightningModule):
    def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
        super().__init__()
        self.input_size = input_size
        self.xcol = xcol
        self.ycol = ycol
        self.layers = nn.Sequential(
            nn.Linear(self.input_size, 1024),
            #nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(1024, 128),
            #nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            #nn.ReLU(),
            nn.Dropout(0.1),

            nn.Linear(64, 16),
            #nn.ReLU(),

            nn.Linear(16, 1)
        )

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

    def training_step(self, batch, batch_idx):
            x = batch[self.xcol]
            y = batch[self.ycol].reshape(-1, 1)
            x_hat = self.layers(x)
            loss = F.mse_loss(x_hat, y)
            return loss
    
    def validation_step(self, batch, batch_idx):
        x = batch[self.xcol]
        y = batch[self.ycol].reshape(-1, 1)
        x_hat = self.layers(x)
        loss = F.mse_loss(x_hat, y)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

def normalized(a, axis=-1, order=2):
    # import numpy as np  # pylint: disable=import-outside-toplevel

    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
    l2[l2 == 0] = 1
    return a / np.expand_dims(l2, axis)


model = MLP(768)  # CLIP embedding dim is 768 for CLIP ViT L 14
model_path = "waifu-scorer-v3.pth"
if torch.cuda.is_available():
    s = torch.load(model_path)   # load the model you trained previously or the model available in this repo
else:
    s = torch.load((model_path), map_location=torch.device('cpu'))

model.load_state_dict(s)
model.to("cuda")
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model2, preprocess = clip.load("ViT-L/14", device=device)  #RN50x64   


pil_image = Image.open(img_path)

image = preprocess(pil_image).unsqueeze(0).to(device)



with torch.no_grad():
   image_features = model2.encode_image(image)

im_emb_arr = normalized(image_features.cpu().detach().numpy() )

prediction = model(torch.from_numpy(im_emb_arr).to(device).type(torch.cuda.FloatTensor))

print( "Aesthetic score predicted by the model:")
print( prediction )

試したこと・調べたこと

• teratailやGoogle等で検索した
• ソースコードを自分なりに変更した
• 知人に聞いた
• その他

上記の詳細・結果

上記のモデルはパラメータのみ保存されており、使用するためには自分でクラスを作成する必要があると認識しています。
しかしながら、どのような層の構造にすれば良いかが分かりません。

補足

モデルのOrderedDictは次のようになっています。
layers.0.weight torch.Size([2048, 768])
layers.0.bias torch.Size([2048])
layers.2.weight torch.Size([2048])
layers.2.bias torch.Size([2048])
layers.2.running_mean torch.Size([2048])
layers.2.running_var torch.Size([2048])
layers.2.num_batches_tracked torch.Size([])
layers.4.weight torch.Size([512, 2048])
layers.4.bias torch.Size([512])
layers.6.weight torch.Size([512])
layers.6.bias torch.Size([512])
layers.6.running_mean torch.Size([512])
layers.6.running_var torch.Size([512])
layers.6.num_batches_tracked torch.Size([])
layers.8.weight torch.Size([256, 512])
layers.8.bias torch.Size([256])
layers.10.weight torch.Size([256])
layers.10.bias torch.Size([256])
layers.10.running_mean torch.Size([256])
layers.10.running_var torch.Size([256])
layers.10.num_batches_tracked torch.Size([])
layers.12.weight torch.Size([128, 256])
layers.12.bias torch.Size([128])
layers.14.weight torch.Size([128])
layers.14.bias torch.Size([128])
layers.14.running_mean torch.Size([128])
layers.14.running_var torch.Size([128])
layers.14.num_batches_tracked torch.Size([])
layers.16.weight torch.Size([32, 128])
layers.16.bias torch.Size([32])
layers.18.weight torch.Size([1, 32])
layers.18.bias torch.Size([1])

使用したプログラム

import torch

state_dict = torch.load("waifu-scorer-v3.pth")

# OrderedDict の内容を確認
for param_tensor in state_dict:
    print(param_tensor, "\t", state_dict[param_tensor].size())

Comments

[meg_]

> こちらのソースコードを使用しました（モデルの説明文に記載されていたため）
> https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/simple_inference.py
"Inspired by"と書かれているので影響を受けて開発したという意味でモデルは別な気がしますが、他に情報がないですね・・・

Answer 1

以下の層構造にしたうえで、state_dictのキーを変更することで実行できました。

ただし実行結果はリンク先のサンプルとは一致せず、0～10の範囲外になる時もあります。
しかし数値の大小は一致しており、他の画像で試した際も直感通りの結果が得られました。

個人的なプログラムで使用する分には上記の結果で十分なため解決済みとさせて頂きます。

# モデルの層構造
nn.Linear(768, 2048), 
nn.BatchNorm1d(2048),    
nn.Tanh(),
nn.Linear(2048, 512),   
nn.BatchNorm1d(512),   
nn.Tanh(),
nn.Linear(512, 256), 
nn.BatchNorm1d(256),
nn.Tanh(),
nn.Linear(256, 128),   
nn.BatchNorm1d(128),       
nn.Tanh(),
nn.Linear(128, 32),      
nn.Linear(32, 1)     

model = MLP(768)
model_path = "waifu-scorer-v3.pth"
s = torch.load(model_path) 

# state_dictのキーを変更
from collections import OrderedDict
new_state_dict = OrderedDict()
for key, value in s.items():
    new_key = key.replace("layers.2", "layers.1")
    new_key = new_key.replace("layers.4", "layers.3")
    new_key = new_key.replace("layers.6", "layers.4")
    new_key = new_key.replace("layers.8", "layers.6")
    new_key = new_key.replace("layers.10", "layers.7")
    new_key = new_key.replace("layers.12", "layers.9")
    new_key = new_key.replace("layers.14", "layers.10")
    new_key = new_key.replace("layers.16", "layers.12")
    new_key = new_key.replace("layers.18", "layers.13")
    new_state_dict[new_key] = value
model.load_state_dict(new_state_dict)