ImageFolderで自作データセットを作成し、PyTorchでCNNをしたい

Question

前提・実現したいこと

ImageFolderを使って自作データセットを作成し、PyTorchでCNNを行いたいと思うのですが、エラーが出てしまいます。

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

Traceback (most recent call last):
  File "/Users/○○○/Desktop/△△△_data/cnn.py", line 125, in <module>
    train_loss_list, test_loss_list = run(30, optimizer, criterion, device)
  File "/Users/○○○/Desktop/△△△_data/cnn.py", line 73, in run
    train_loss = train_epoch(model, optimizer, criterion, train_loader, device)
  File "/Users/○○○/Desktop/△△△_data/cnn.py", line 46, in train_epoch
    outputs = model(images)
  File "/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/○○○/Desktop/△△△_data/cnn.py", line 33, in forward
    x = x.view(-1, 16*5*5)
RuntimeError: shape '[-1, 400]' is invalid for input of size 2080832

該当のソースコード

from pathlib import Path

from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image

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

import matplotlib.pyplot as plt

from torchvision.datasets import ImageFolder

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 畳み込み層:(入力チャンネル数, フィルタ数、フィルタサイズ)
        # 出力画像サイズ28
        self.pool = nn.MaxPool2d(2, 2)  # プーリング層:（領域のサイズ, ストライド）
        # 出力画像サイズ14
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 出力画像サイズ5
        self.fc1 = nn.Linear(16*5*5, 256)  # 全結合層
        self.dropout = nn.Dropout(p=0.5)  # ドロップアウト:(p=ドロップアウト率)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16*5*5)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x


def train_epoch(model, optimizer, criterion, dataloader, device):
    train_loss = 0
    model.train()
    for i, (images, labels) in enumerate(dataloader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    train_loss = train_loss / len(train_loader.dataset)
    return train_loss


def inference(model, optimizer, criterion, dataloader, device):
    model.eval()
    test_loss=0

    with torch.no_grad():
        for i, (images, labels) in enumerate(test_loader):
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            test_loss += loss.item()
        test_loss = test_loss / len(test_loader.dataset)
    return test_loss


def run(num_epochs, optimizer, criterion, device):
    train_loss_list = []
    test_loss_list = []
    for epoch in range(num_epochs):
        train_loss = train_epoch(model, optimizer, criterion, train_loader, device)
        test_loss = inference(model, optimizer, criterion, test_loader, device)

        print(f'Epoch [{epoch+1}], train_Loss : {train_loss:.4f}, test_Loss : {test_loss:.4f}')
        train_loss_list.append(train_loss)
        test_loss_list.append(test_loss)
    return train_loss_list, test_loss_list


if __name__ == '__main__':
    # Transform を作成する。
    transform = transforms.Compose([transforms.Resize(256), transforms.ToTensor()])
    # Dataset を作成する。
    dataset = ImageFolder('/Users/○○○/Desktop/△△△_data/□□□_dataset', transform)
    # DataLoader を作成する。
    dataloader = DataLoader(dataset, batch_size=3)

    #for batch in dataloader:
    #   print(batch.shape)

    # グラフのスタイルを指定
    plt.style.use('seaborn-darkgrid')

    # 正規化
    normalize = transforms.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0))  
    #  Tensor化
    to_tensor = transforms.ToTensor()

    train_ratio = 0.8
    train_size = int(train_ratio * len(dataset))
    val_size  = len(dataset) - train_size      
    data_size  = {"train":train_size, "val":val_size}
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

    transform_train = transforms.Compose([to_tensor, normalize])
    transform_test = transforms.Compose([to_tensor, normalize])

    batch_size = 64
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader   = torch.utils.data.DataLoader(test_dataset,   batch_size=batch_size, shuffle=False)
    dataloaders  = {"train":train_loader, "val":test_loader}

    model = Net()

    device = torch.device("cpu")

    # 損失関数の設定
    criterion = nn.CrossEntropyLoss()

    # 最適化手法を設定
    optimizer = optim.Adam(model.parameters())

    train_loss_list, test_loss_list = run(30, optimizer, criterion, device)

    num_epochs=30
    fig, ax = plt.subplots(figsize=(8, 6), dpi=100)
    ax.plot(range(num_epochs), train_loss_list, c='b', label='train loss')
    ax.plot(range(num_epochs), test_loss_list, c='r', label='test loss')
    ax.set_xlabel('epoch', fontsize='20')
    ax.set_ylabel('loss', fontsize='20')
    ax.set_title('training and test loss', fontsize='20')
    ax.grid()
    ax.legend(fontsize='25')
    plt.show()

    correct = 0
    total = 0
    model.eval()  # 評価モード
    for i, (images, label) in enumerate(test_loader):
        images, label = images.to(device), label.to(device) 
        prediction = model(images)
        correct += (prediction.argmax(1) == label).sum().item()
        total += len(images)
    print(f"正解率: {(correct/total)*100:.3f} %")

該当のデータセット

./□□□_dataset
├── 0
│   ├── 0_0.png
│   ├── 0_1.png
│   ├── ・・・
│   ├── 0_10.png
├── 1
│   ├── 1_0.png
│   ├── 1_1.png
│   ├── ・・・
│   ├── 1_10.png
├── 2
│   ├── 2_0.png
│   ├── 2_1.png
│   ├── ・・・
│   ├── 2_10.png

Comments

[tiitoi]

パスは合っていますか？
画像パスの一覧が取得できていれば、dataset.img_paths の値を print したときに中身があるはずですがどうでしょうか？

[k_0fall]

print(dataset.img_paths)で確認したところ、
[]
という出力になってしまい、正しく画像を読み込めていなかったです。

内容としては、
./△△△△△_dataset
├── 0
│ ├── 0_0.png
│ ├── 0_1.png
│ ├── ・・・
│ ├── 0_10.png
├── 1
│ ├── 1_0.png
│ ├── 1_1.png
│ ├── ・・・
│ ├── 1_10.png
├── 2
│ ├── 2_0.png
│ ├── 2_1.png
│ ├── ・・・
│ ├── 2_10.png
という構造にして、合計30枚の画像を読み込む予定です。

Answer 1

./△△△△△_dataset
├── 0
│ ├── 0_0.png
│ ├── 0_1.png
│ ├── ・・・
│ ├── 0_10.png
├── 1
│ ├── 1_0.png
│ ├── 1_1.png
│ ├── ・・・
│ ├── 1_10.png
├── 2
│ ├── 2_0.png
│ ├── 2_1.png
│ ├── ・・・
│ ├── 2_10.png

このような構造であれば、ImageFolder を自作する必要はないです。
torchvision.datasets.ImageFolder をお使いください。

from pathlib import Path

from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.datasets import ImageFolder

# Transform を作成する。
transform = transforms.Compose([transforms.Resize(256), transforms.ToTensor()])
# Dataset を作成する。
dataset = ImageFolder("dataset1", transform)

Pytorch - 自作のデータセットを扱う Dataset クラスを作る方法

追記

必要なサイズがどのような数値が判断する方法などありましたらご教授頂けないでしょうか..

全結合層の直前までのモデルを作成し、希望のテンソルを入れて出力の形状を確認すればよいです。

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        return x


model = Net()

x = torch.rand(1, 3, 224, 224)
y = model(x)
print(y.shape)  # torch.Size([1, 16, 53, 53])

Comments

[k_0fall]

ご回答ありがとうございます。
ImageFolderではない形にしても、同様のエラーが発生します。

データセットからラベルを無くし、以下のようにすべて同じディレクトリに配置したところ、別のエラーになりました。
./△△△△△_dataset
├── 0_0.png
├── 0_1.png
├── ・・・
├── 0_10.png
├── 1_0.png
├── 1_1.png
├── ・・・
├── 1_10.png
├── 2_0.png
├── 2_1.png
├── ・・・
├── 2_10.png

torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([3, 4, 256, 341])
torch.Size([2, 4, 256, 341])
Traceback (most recent call last):
File "/Users/○○○○○/Desktop/●●_data/cnn.py", line 180, in <module>
train_loss_list, test_loss_list = run(30, optimizer, criterion, device)
File "/Users/○○○○○/Desktop/●●_data/cnn.py", line 111, in run
train_loss = train_epoch(model, optimizer, criterion, train_loader, device)
File "/Users/○○○○○/Desktop/●●_data/cnn.py", line 81, in train_epoch
for i, (images, labels) in enumerate(dataloader):
ValueError: too many values to unpack (expected 2)

[tiitoi]

元のディレクトリ構造で、from torchvision.datasets import ImageFolder を使ってもエラーになりましたか？ (自作したImageFolder ではなく)

[k_0fall]

すみません、一度from torchvision.datasets import ImageFolderを使いましたが、エラーが出ています。
複雑になってしまったので、質問を更新しました。
エラー、ソースコード、ディレクトリは更新した通りの内容です。

[tiitoi]

データセット読み込みのところは正常に動いています。
質問のコードはどこから持ってきたものなのでしょうか？
入力画像サイズをそのモデルで想定されている値にリサイズする必要があります。

↓ にリサイズ処理を追加する必要がある
transform_train = transforms.Compose([to_tensor, normalize])
transform_test = transforms.Compose([to_tensor, normalize])

[k_0fall]

CNNの学習については主に、
https://dreamer-uma.com/pytorch-cnn/
自作データセットにおけるtrainとtestの振り分け？のような動作には、
https://qiita.com/ryryry/items/b1da4855504dcd3f9d98
のサイトを参考にしました。

リサイズ処理について詳しく調べてみます。

[tiitoi]

transforms.RandomResizedCrop() や transforms.Resize() を使ってリサイズします

https://pystyle.info/pytorch-train-classification-problem-using-a-pretrained-model/#outline__4_3

[k_0fall]

ありがとうございます！
リサイズを試してみました。

その結果、エラーの1番下の「RuntimeError: shape '[-1, 400]' is invalid for input of size」の後の数字が2080832から、1213488に変わりました。

このエラーについて調べたところ、全結合層のサイズが関係しているそうです。
参考にしたソースコードはtorchvision.datasets.CIFAR10を想定して16*5*5というサイズだそうですが、自作のデータセットの場合は違うため、
self.fc1 = nn.Linear(16*5*5, 256) や
x = x.view(-1, 16*5*5)の箇所を変える必要があるそうです。

必要なサイズがどのような数値が判断する方法などありましたらご教授頂けないでしょうか..

[tiitoi]

全結合層の直前までのモデルを作成して希望の入力サイズのテンソルを入れてみて、出力の形状を確認するとよいです。
回答に追記しました。

Answer 2

ありがとうございます！
165353にすることで解決することができました。