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

前提

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種類試してみる

python3 〜〜〜〜.py --prompt "a photograph of an astronaut riding a horse" --plms
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

uuuuuu

2023/01/18 10:16

バッチサイズを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