PytorchのMNISTでエラーが出ます

Question

前提・実現したいこと

Pytorchで機械学習をしたところエラーが出ます。
Pytorchはhttps://www.tensorflow.org/tutorials/quickstart/advanced?hl=ja
こちらのTensorflowの実装と同じ実装をしたいと思っています。

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

ValueError: Expected input batch_size (784) to match target batch_size (32).

該当のソースコード

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F


transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, ), (0.5, ))])
trainset = torchvision.datasets.MNIST(root='./data', 
                                        train=True,
                                        download=True,
                                        transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
                                            batch_size=32,
                                            shuffle=True,
                                            num_workers=2)

testset = torchvision.datasets.MNIST(root='./data', 
                                        train=False, 
                                        download=True, 
                                        transform=transform)
testloader = torch.utils.data.DataLoader(testset, 
                                            batch_size=32,
                                            shuffle=False, 
                                            num_workers=2)

print("trainset.shape:",trainset)
print('train_dataset = ', len(trainset))  
print('test_dataset = ', len(testset))

classes = tuple(np.linspace(0, 9, 10, dtype=np.uint8))


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)  #(入力のchannel数, 出力のチャンネル数, カーネル)
        self.flatten = nn.Flatten()
        self.d1 = nn.Linear(32, 128)
        self.d2 = nn.Linear(128, 10)
        
    def forward(self, x):
        print(x.shape)
        #torch.reshape(x, (32, 32))
        x = x.view(-1, 32)
        x = self.flatten(x)
        x = F.relu(self.d1(x))
        return F.softmax(self.d2(x))
 
# select device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = Net().to(device)
 
# optimizing
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

epochs = 5
for epoch in range(epochs):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(trainloader, 0):
        # zero the parameter gradients
        #inputs, labels = inputs.view(-1, 28*28*1).to(device), labels.to(device)
        optimizer.zero_grad()
        #print(input)

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 100 == 99:
            print('[{:d}, {:5d}] loss: {.3f}'
                    .format(epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

print('Finished Training')

試したこと

x = x.view(-1, 32)を付け加えました。

補足情報（FW/ツールのバージョンなど）

batchサイズをtensorflowの方と同じにしたいです。

Answer 1

testloader = torch.utils.data.DataLoader(testset, 
                                            batch_size=32,`
                                            shuffle=False, 
                                            num_workers=2)

で指定してあるので，forward()の入力テンソルのサイズは(32, 1, 28, 28)となっており，view()等は不要です．
実際のミスは別のところに3つあります．

1. self.conv1(x)を呼び忘れている
2. conv1のあとのReLUを呼び忘れている
3. conv1の出力は，(32, 32, 26, 26)なので，d1の入力チャネル数は32 * 26 * 26になる

それらを直すと以下のようになります．

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)  #(入力のchannel数, 出力のチャンネル数, カーネル)
        self.flatten = nn.Flatten()
        self.d1 = nn.Linear(32 * 26 * 26, 128)
        self.d2 = nn.Linear(128, 10)

    def forward(self, x):
        print(x.shape)
        x = F.relu(self.conv1(x))
        x = self.flatten(x)
        x = F.relu(self.d1(x))
        return F.softmax(self.d2(x))

さらにto(device)もしておけば動くと思います．

- #inputs, labels = inputs.view(-1, 28*28*1).to(device), labels.to(device)
+ inputs, labels = inputs.to(device), labels.to(device)

Comments

[SaitoHiroaki]

丁寧な回答ありがとうございます。
ご指摘いただいた点を改善してみたところ、新たなエラーがでました。

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

というエラーが出ました。
この改善策もアドバイスをいただけないでしょうか

[退会済みユーザー]

weightがGPUメモリに乗っていないというエラーなので，`net.to(device)` が漏れているような気がします．
解決しなければ，質問を別に立ててください（質問を増やすのを繰り返されても困るので）．

[SaitoHiroaki]

わかりました。ありがとうございます。