pytorchを用いてCNNによる2種類の画像分類を行いたいが.stack expects each tensor to be equal sizeというエラーが出てしまう

Question

前提

Linux
GPU
python3
pytorch

画像データは正方形、サイズ：255＊255に揃えたものを使用しています。

実現したいこと

python3とpytorchを用いてCNNによるYESかNOの画像分類を行いたいです。

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

画像の枚数を100枚：50枚→400枚：100枚（学習用：テスト用）に変更したところ以下のエラーが発生し実行ができなくなりました。
画像を増やす前は実行ができていました。

**CUDA out of memory. **Tried to allocate 352.00 MiB (GPU 0; 7.92 GiB total capacity;
7.37 GiB already allocated; 107.69 MiB free; 7.43 GiB reserved in total by PyTorch)
 If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

該当のソースコード

#CNN学習　#memo7_GPUより
import copy
import time
import os
import tqdm

import torchvision.transforms as transforms
import torchvision.models as models
import torchvision
from torchvision import transforms
from torch.utils.data import (DataLoader,Dataset,TensorDataset)
import torch.optim as optim
import torch.nn as nn
import torch

def saveModel(): 
    path="./sim.pth "  

#modelを保存するファイル名
    torch.save(model.state_dict(),path)


data_tr_folder = "./learning98"#data写真データのフォルダ（学習）
data_vl_folder = "./test1" #data写真データのフォルダ（テスト）


transform_dict = {
    
    'train': transforms.Compose([
    #transforms.RandomCrop(127),#画像リサイズいろいろ。ramdomarop以外にもあるから調べてみる
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ]),
    'val': transforms.Compose([
    #transforms.CenterCrop(127),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
])}




transform   = transform_dict["train"]

data_tr = torchvision.datasets.ImageFolder(root=data_tr_folder, transform=transform_dict["train"])
data_vl = torchvision.datasets.ImageFolder(root=data_vl_folder, transform=transform_dict["val"])



train_size = int(len(data_tr))
# int()で整数に。
val_size  = int(len(data_vl))
data_size  = {"train":train_size, "val":val_size}
#          =>{"train": 4800,       "val":800}
#data_train = torch.utils.data.random_split(data_tr,train_size)
#data_val = torch.utils.data.random_split(data_vl,val_size)
batch_size=32   #バッチサイズ

train_loader = torch.utils.data.DataLoader(data_tr, batch_size=batch_size, shuffle=True)
#val_loaderをtest_loaderに名前かえてみた
test_loader   = torch.utils.data.DataLoader(data_vl, batch_size=batch_size, shuffle=True)
dataloaders  = {"train":train_loader, "val":test_loader}


#上記データセット作成

# リスト4.2　2層の畳み込み層と2層のMLPをつなげたCNNを作成

# In[ ]:


#  (N、C、H、W)形式のTensorを(N, C*H*W)に引き伸ばす層
# 畳み込み層の出力をMLPに渡す際に必要
class FlattenLayer(nn.Module):
    def forward(self, x):
        sizes = x.size()
        return x.view(sizes[0], -1)
       # size = x.size()[1:]  # all dimensions except the batch dimension
        #num_features = 1
        #for s in size:
        #    num_features *= s
       # return num_features





#CNNの本体

conv_net = nn.Sequential(
    nn.Conv2d(3, 32, 5),#カラーなら３でOK、３２は画像のサイズによって変更したらうまくいくかも
    nn.MaxPool2d(2),
    nn.ReLU(),
    nn.BatchNorm2d(32),
    nn.Dropout2d(0.5),    ###
    nn.Conv2d(32, 64, 5),
    nn.MaxPool2d(2),
    nn.ReLU(),
    nn.BatchNorm2d(64),
    nn.Dropout2d(0.5),    ###
    nn.Conv2d(64, 128, 5),
    nn.MaxPool2d(2),
    nn.ReLU(),
    nn.BatchNorm2d(128),
    FlattenLayer()
)


# 畳み込みによって最終的にどのようなサイズになっているかを、
# 実際にダミーデータを入れてみて確認する
test_input = torch.ones(1, 3, 227, 227)  #RGBか。画像のサイズ。
conv_output_size = conv_net(test_input).size()[-1]

# 2層のMLP
mlp = nn.Sequential(
    nn.Linear(conv_output_size, 200),
    nn.ReLU(),
    nn.BatchNorm1d(200),
    nn.Dropout(0.5),
    nn.Linear(200, 10)
)

# 最終的なCNN
net = nn.Sequential(
    conv_net,
    mlp
)


# リスト4.3　評価と訓練のヘルパー関数を作成


# In[ ]:


# 評価のヘルパー関数
def eval_net(net, data_loader, device="cuda:0"):
    # DropoutやBatchNormを無効化
    net.eval()
    ys = []
    ypreds = []
    for x, y in data_loader:
        # toメソッドで計算を実行するデバイスに転送する
        x = x.to(device)
        y = y.to(device)
        # 確率が最大のクラスを予測(リスト2.14参照)
        # ここではforward（推論）の計算だけなので自動微分に
        # 必要な処理はoffにして余計な計算を省く
        with torch.no_grad():
            _, y_pred = net(x).max(1)
        ys.append(y)
        ypreds.append(y_pred)
            # ミニバッチごとの予測結果などを1つにまとめる
    ys = torch.cat(ys)
    ypreds = torch.cat(ypreds)
    # 予測精度を計算
    acc = (ys == ypreds).float().sum() / len(ys)
    return acc.item()
  
# 訓練のヘルパー関数
def train_net(net, train_loader, test_loader,
              optimizer_cls=optim.Adam,
              loss_fn=nn.CrossEntropyLoss(),
              n_iter=10, device="cuda:0"):
    best_accuracy = 1.0 #save?#
    train_losses = []
    train_acc = []
    val_acc = []
    optimizer = optimizer_cls(net.parameters(),lr=0.00001)  ###
    for epoch in range(n_iter):
        running_loss = 0.0
        # ネットワークを訓練モードにする
        net.train()
        n = 0
        n_acc = 0
        # 非常に時間がかかるのでtqdmを使用してプログレスバーを出す
        for i, (xx, yy) in tqdm.tqdm(enumerate(train_loader),
            total=len(train_loader)):
            xx = xx.to(device)
            yy = yy.to(device)
            h = net(xx)
            loss = loss_fn(h, yy)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            n += len(xx)
            _, y_pred = h.max(1)
            n_acc += (yy == y_pred).float().sum().item()
        train_losses.append(running_loss / i)
        # 訓練データの予測精度
        train_acc.append(n_acc / n)
        # 検証データの予測精度
        val_acc.append(eval_net(net, test_loader, device))

       # accuracy = train_losses[-1]
       # if accuracy < best_accuracy:
       #     saveModel()
       #     best_accuracy=accuracy
            

            # このepochでの結果を表示
        print(epoch, train_losses[-1], train_acc[-1],
              val_acc[-1], flush=True)
        


# リスト4.4　全パラメータをGPUに転送して訓練を実行

# In[ ]:


# ネットワークの全パラメータをGPUに転送
model=net.to("cuda:0")
# 訓練を実行
train_net(model, train_loader, test_loader, n_iter=150, device="cuda:0")
saveModel()

試したこと

pytorch GPU メモリ開放　or メモリ不足と調べて出てきた記事に書いてあることを読み漁り実践してみましたがどれもうまく行きませんでした。

具体的には2つのことを行いました

• リストメモリの量を確認してkillしてみる
• リスト実行コマンドを2種類試してみる

1. python3 〜〜〜〜.py --prompt "a photograph of an astronaut riding a horse" --plms
2. PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:100 python3 〜〜〜〜.py

メモリの確認画像

追加記載

バッチサイズを64→32に変更したところエラーが消えました。
が、違うエラーが発生しました。
RuntimeError: stack expects each tensor to be equal size, but got [3, 512, 512] at entry 0 and [3, 1095, 1280] at entry 22

Comments

[uuuuuu]

バッチサイズを64→32に変更したところエラーが消えました。
が、違うエラーが発生しました。
RuntimeError: stack expects each tensor to be equal size, but got [3, 512, 512] at entry 0 and [3, 1095, 1280] at entry 22

[meg_]

out of memoryの問題は解決したということでしょうか？ であれば質問タイトルから変更された方が良いのではないでしょうか？
「違うエラー」はコードのどこで発生したのか分かりますか？

[uuuuuu]

コメントありがとうございます。
おそらくコード内ではなく画像データのサイズが合ってないというエラーなのですが、
画像サイズは512*512のjpegに全て揃えており、枚数も学習用とテスト用のYES、NOでそれぞれ合わせているので原因がわからず困っている事態です。

[meg_]

> 画像サイズは512*512のjpegに全て揃えており、
それでも[3, 1095, 1280]のエラーが出るということは、サイズ違いの画像が紛れている可能性が高いのでしょうね。
枚数が多く目視での確認が困難であれば、プログラムで画像サイズを確認されるのが良いかと思います。

[uuuuuu]

512＊512で揃っているかの確認としてもう一度一括で揃えるコマンド
mogrify -gravity center -resize 512x512 -extent 512x512 *.jpeg

を実行したので紛れている可能性もないはずなんです...困り果てました笑

[meg_]

identify -format '%wx%h\n' *
などで確認されてはどうでしょうか？

[uuuuuu]

気持ちいいくらいに解決できました！！！！！
一つjpegに揃っていないものがいました。
大変助かりました。最高です。

[uuuuuu]

ぜひ解答欄に記入していただけないでしょうか？