Tensor for argument #2 'mat1' is on CPU, but expected it to be on GPU の対処方法

Question

1.前提・実現したいこと
pytorchによる変分オートエンコーダの実装です。
google　colab上で動かしています。
下記（2.）にエラーメッセージを示しました。
参考コードにある

784次元ベクトルを2次元ベクトルにencode

z = model.encode(Variable(images, volatile=True))
の部分でエラーが発生します。
gpu、cpuがらみのエラーなのでしょうか？
エラー解消法のご教授をお願いいたします。

2.発生している問題・エラーメッセージ
RuntimeError: Tensor for argument #2 'mat1' is on CPU, but expected it to be on GPU (while checking arguments for addmm)

3.参考コード
引用元は、「人工知能に関する断創録」様の下記のブログ（更新停止: 2019年12月31日）です。
https://aidiary.hatenablog.com/entry/20180228/1519828344

エポックは、時間短縮で10にしていますが、100でも動くことは確認済みです。

import os
import numpy as np
import torch
import torch.nn as nn
import torch.utils.data
import torch.optim as optim
from torch.autograd import Variable
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image

batch_size = 128
num_epochs = 10
seed = 1
out_dir = './vae_2'

cuda = torch.cuda.is_available()
if cuda:
    print('cuda is available!')

if not os.path.exists(out_dir):
    os.mkdir(out_dir)

torch.manual_seed(seed)
if cuda:
    torch.cuda.manual_seed(seed)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data',
                   train=True,
                   download=True,
                   transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data',
                   train=False,
                   transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=True
)

class VAE(nn.Module):

    def __init__(self):
        super(VAE, self).__init__()
        
        self.fc1 = nn.Linear(784, 512)
        self.fc21 = nn.Linear(512, 2)  # mu
        self.fc22 = nn.Linear(512, 2)  # logvar

        self.fc3 = nn.Linear(2, 512)
        self.fc4 = nn.Linear(512, 784)
        
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
    
    def encode(self, x):
        h = self.relu(self.fc1(x))
        return self.fc21(h), self.fc22(h)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.data.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h = self.relu(self.fc3(z))
        return self.sigmoid(self.fc4(h))
    
    def forward(self, x):
        x = x.view(-1, 784)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

model = VAE()
if cuda:
    model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

def loss_function(recon_x, x, mu, logvar):
    # size_average=Falseなのでバッチ内のサンプルの合計lossを求める
    # reconstruction loss 入力画像をどのくらい正確に復元できたか？
    # 数式では対数尤度の最大化だが交差エントロピーlossの最小化と等価
    recon = F.binary_cross_entropy(recon_x, x.view(-1, 784), size_average=False)

    # 潜在空間zに対する正則化項
    # P(z|x) が N(0, I)に近くなる（KL-distanceが小さくなる）ようにする
    kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon + kld

#https://blog.tamesuu.com/2019/09/24/415/ error対応
def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        if cuda:
            data = Variable(data.cuda())
        else:
            data = Variable(data)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.data.item()#error対応
        optimizer.step()
    
    # loss_function() は平均ではなく全サンプルの合計lossを返すのでサンプル数で割る
    train_loss /= len(train_loader.dataset)

    return train_loss    
    

def test(epoch):
    model.eval()
    test_loss = 0
    for batch_idx, (data, _) in enumerate(test_loader):
        if cuda:
            data = Variable(data.cuda(), volatile=True)
        else:
            data = Variable(data, volatile=True)
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        test_loss += loss.data.item()#error対応
        
        if epoch % 10 == 0:
            # 10エポックごとに最初のminibatchの入力画像と復元画像を保存
            if batch_idx == 0:
                n = 8
                comparison = torch.cat([data[:n],
                                        recon_batch.view(batch_size, 1, 28, 28)[:n]])
                save_image(comparison.data.cpu(),
                           '{}/reconstruction_{}.png'.format(out_dir, epoch), nrow=n)

    test_loss /= len(test_loader.dataset)

    return test_loss

loss_list = []
test_loss_list = []
for epoch in range(1, num_epochs + 1):
    loss = train(epoch)
    test_loss = test(epoch)

    print('epoch [{}/{}], loss: {:.4f} test_loss: {:.4f}'.format(
        epoch + 1,
        num_epochs,
        loss,
        test_loss))

    # logging
    loss_list.append(loss)
    test_loss_list.append(test_loss)

# save the training model
np.save('loss_list.npy', np.array(loss_list))
np.save('test_loss_list.npy', np.array(test_loss_list))
torch.save(model.state_dict(), 'vae.pth')

# Commented out IPython magic to ensure Python compatibility.
import matplotlib.pyplot as plt
# %matplotlib inline
loss_list = np.load('/content/loss_list.npy'.format(out_dir))

plt.plot(loss_list)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.grid()

test_loss_list = np.load('/content/test_loss_list.npy'.format(out_dir))
plt.plot(test_loss_list)
plt.xlabel('epoch')
plt.ylabel('test loss')
plt.grid()

from IPython.display import Image
Image('vae_2/reconstruction_10.png')

Image('vae_2/reconstruction_100.png')

model.load_state_dict(torch.load('/content/vae.pth'.format(out_dir),
                                 map_location=lambda storage,
                                 loc: storage))
test_dataset = datasets.MNIST('./data', download=True, train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=10000, shuffle=False)
images, labels = iter(test_loader).next()
images = images.view(10000, -1)

# 784次元ベクトルを2次元ベクトルにencode
z = model.encode(Variable(images, volatile=True))

mu, logvar = z
mu, logvar = mu.data.numpy(), logvar.data.numpy()
print(mu.shape, logvar.shape)

import pylab
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 10))
plt.scatter(mu[:, 0], mu[:, 1], marker='.', c=labels.numpy(), cmap=pylab.cm.jet)
plt.colorbar()
plt.xlim((-6, 6))
plt.ylim((-6, 6))
plt.grid()

4.補足情報（FW/ツールのバージョンなど）
開発環境：Google Colaboratory
プログラム言語：python3
OS：windows10 Home
CPU：Intel(R) Core(TM) i7-7500U CPU@2.70GHz 2.90GHz

Answer 1

z = model.encode(Variable(images, volatile=True))
の部分でエラーが発生します。
gpu、cpuがらみのエラーなのでしょうか？

基本的にはそういうことだと思います。
参考: pytorch でハマりやすいポイント
https://ryuiso2.hatenablog.com/entry/2020/01/13/013858

エラー解消法のご教授をお願いいたします。

images というのがまだ GPU側に行っていないことが基本的な問題なんだと思います。

z = model.encode(Variable(images, volatile=True).cuda())

とか、 .to(<<device>>) みたいな感じでどうでしょうか。

Comments

[kouji_39]

回答者さまのご指導どおりに修正してみました。
mu, logvar = mu.data.numpy(), logvar.data.numpy()のところで、
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.というエラーがでます。

# 784次元ベクトルを2次元ベクトルにencode
z = model.encode(Variable(images, volatile=True).cuda())
#z = model.encode(Variable(images, volatile=True))

mu, logvar = z
＃下記コードでエラー発生
mu, logvar = mu.data.numpy(), logvar.data.numpy()
print(mu.shape, logvar.shape)

[mokemokechicken]

今度は、 Tensor が GPU 側(cuda:0)にあるので、CPU側に持ってこれないので

mu, logvar = mu.data.cpu().numpy(), logvar.data.cpu().numpy()

みたいな感じでどうでしょうか。

[kouji_39]

mokemokechickenさん、適切なご指導ありがとうございます。
うまく作動しました。