[Python] Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the sameの原因と解決方法

2022年5月27日金曜日

Python

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensorのエラー原因と解決方法を紹介します。

アイキャッチ

結論

モデルと入力データを配置しているデバイスが異なるため発生しています。
上記の場合モデルをGPUに配置し、入力データをCPUに配置しているため発生しています。CPUかGPUどちらかで統一する必要があります。

下記エラー発生例です。

!pip install transformers > /dev/null

from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import requests

url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)

feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
# ロードしたモデルをGPUに配置
model.to('cuda:0')

inputs = feature_extractor(images=image, return_tensors="pt")
# 入力データをCPUに配置
inputs.to('cpu')

# ここでエラー発生
# RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor
outputs = model(**inputs)

モデルと入力データともに'cpu'または'cuda:0'に揃えてあげることで解消します。

よくある原因

よくある原因は以下の通りです。

  1. デフォルトで展開されたデバイスが異なっている
    先ほどのケースだと明示的に配置先を指定していますが、指定しなかった場合にデバイスの配置先が異なってしまうケースです。
  2. メモリを解放しきれていない
    事前に、何らかの処理をしており一度配置したモデルなどが意図せず残っており上記エラーが発生するケースです。
    ランタイムを一度再起動するなどメモリを解放すると解消する可能性があります。

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology