Google Colaboratoryを使った機械学習でイメージのサイズが合わない

実現したいこと

機械学習を行ってモデルを保存したい

前提

学校の授業でpythonを用いて機械学習の実験を行っています。
今までの授業と同じようにコードを書いているはずなのにエラーが出てきました。
友達はこれで出来たと言っていたのですが...

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

Traceback (most recent call last):
  File "src/train_for_beautiful.py", line 121, in <module>
    main()  # main()が実行される
  File "src/train_for_beautiful.py", line 95, in main
    train(net, trainloader, optimizer, criterion, args.nepochs)
  File "src/train_for_beautiful.py", line 31, in train
    loss = criterion(outputs, labels)  # ロスの計算 
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 3026, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: 0D or 1D target tensor expected, multi-target not supported

該当のソースコード

GoogleColaboratory
1!python src/train_jkj1a.py --nepochs 10 --batch_size 1 --lr 0.01 --save_model_name 'model/model_jkj1a.pth'

train_jkj1a.py

python
1import argparse
2
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6import torch.optim as optim
7
8
9import sys
10sys.path.append('src')
11from vgg_jkj1a import Net
12from jkj1a import load_data
13
14import warnings
15warnings.filterwarnings("ignore")  # warningを表示しない
16
17def train(net, trainloader, optimizer, criterion, nepochs):
18    net.train()  # ネットワークを「訓練モード」にする（おまじない）．
19
20    for epoch in range(nepochs):  
21        # --- ここを埋める ---------------
22        running_loss = 0.0
23        for data in trainloader:
24            images, labels = data
25            images, labels = images.cuda(), labels.cuda()
26            
27            optimizer.zero_grad()
28            
29            outputs = net(images)              # ネットワークへの代入
30            loss = criterion(outputs, labels)  # ロスの計算 
31            loss.backward()                    # 勾配の計算
32            optimizer.step()                   # ネットワークパラメタの更新
33
34            running_loss += loss.item()
35
36        running_loss = running_loss / len(trainloader)
37        print(f'[epoch {epoch + 1:2d}] loss: {running_loss:.3f}')
38        
39        # ------------------------------
40    print('Training completed')
41    
42def test(net, dataloader):
43    net.eval()  # ネットワークを「評価モード」にする（おまじない）．
44
45    correct = 0  # 正解数
46    total = 0    # 画像総数
47
48    for data in dataloader:
49        # --- ここを埋める ---------------
50        images, labels = data
51        outputs = net(images.cuda())
52        predicted = outputs.argmax(dim=1)
53
54        correct += (predicted == labels.cuda()).sum().item()
55        total += len(labels)
56        # ------------------------------
57
58    acc = correct / total
59
60    return acc 
61
62# これはおまじない
63def set_args():
64    # Training settings
65    parser = argparse.ArgumentParser(description='')
66    parser.add_argument('--batch_size', type=int, default=64)
67    parser.add_argument('--nepochs', type=int, default=14)
68    parser.add_argument('--lr', type=float, default=0.01)
69    parser.add_argument('--save_model_name', type=str, default='')
70    args = parser.parse_args()
71    return args 
72
73def main():
74    # --- 変更不要 --------------------------------
75    # GPU or CPUなのか表示
76    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
77    print('using device:', device)
78
79    # `--nepochs 2` などで与えたパラメタの読み込み
80    # `args.nepochs` のようにして使える
81    args = set_args()
82
83    
84    # --- データをロード ---------------------------
85    trainloader, testloader, classes = load_data(args.batch_size)
86
87    
88    
89    
90    # --- ネットワークの初期化と学習------------------
91    net = Net('VGG11')
92    net.cuda()
93    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)  # 確率的勾配法SGD
94    criterion = nn.CrossEntropyLoss()  # ロス関数の定義：分類で標準的なクロスエントロピーロス
95    train(net, trainloader, optimizer, criterion, args.nepochs)
96    net.eval()
97
98    # ------------------------------
99
100    # --- ネットワークを評価して正解率を表示 ---------- 
101    train_acc = test(net, trainloader)
102    test_acc = test(net, testloader)
103    print(f'train acc = {train_acc:.3f}')  # ':.3f'とつけると小数点以下3桁までの表示になる
104    print(f' test acc = {test_acc:.3f}')  
105
106    # --- 変更不要 --------------------------------
107    if args.save_model_name:  # 保存先が与えられている場合保存
108        PATH = args.save_model_name
109        torch.save(net.state_dict(), PATH)
110
111    # state_dict = torch.load(args.save_model_name)        # 保存したパラメタをロード
112    # net.load_state_dict(state_dict)      # ネットワークにパラメタをセット
113    # train_acc = test(net, trainloader)
114    # test_acc = test(net, testloader)
115    # print(f'train acc = {train_acc:.3f}')  # ':.3f'とつけると小数点以下3桁までの表示になる
116    # print(f' test acc = {test_acc:.3f}')  
117
118
119# おまじない（変更しなくて良い）
120if __name__ == '__main__':
121    main()  # main()が実行される

vgg_jkj1a.py

python
1import torch
2import torch.nn as nn
3
4
5cfg = {
6    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
7    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
8    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
9    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
10}
11
12class VGG(nn.Module):
13    def __init__(self, vgg_name):
14        super(VGG, self).__init__()
15        self.features = self._make_layers(cfg[vgg_name])
16        self.classifier = nn.Linear(512, 2)
17
18    def forward(self, x):
19        out = self.features(x)
20        out = out.view(out.size(0), -1)
21        out = self.classifier(out)
22        return out
23
24    def _make_layers(self, cfg):
25        layers = []
26        in_channels = 3
27        for x in cfg:
28            if x == 'M':
29                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
30            else:
31                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
32                           nn.BatchNorm2d(x),
33                           nn.ReLU(inplace=True)]
34                in_channels = x
35        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
36        return nn.Sequential(*layers)

jkj1a.py

python
1from pathlib import Path
2
3from PIL import Image
4
5import os 
6import torch
7import torchvision.transforms as transforms
8from torchvision.datasets import FashionMNIST
9from torch.utils.data import random_split, DataLoader, Subset, Dataset
10import numpy as np
11
12def load_data(batch_size, metric='beauty', img_size=32):  # metric is 'beauty' or 'pleasure'
13
14    # 画像への前処理
15    transform = transforms.Compose(
16            [transforms.Resize(img_size), # original size is 256 x 256
17             transforms.ToTensor(),
18             ] 
19            )
20    
21    # ラベルに対する前処理（各自で適宜修正）
22    target_transform = transforms.Compose(
23        [
24         transforms.Lambda(torch.tensor),
25         transforms.Lambda(lambda x: x / x.norm())  # xは7次元ベクトル（ヒストグラム）．それの正規化している．
26        ]
27    )
28
29    dataset = JKJ1Adataset(
30                transform=transform, 
31                target_transform=target_transform,
32                metric=metric)
33
34    n_test = 50
35    n_train = len(dataset) - n_test
36    trainset, testset = random_split(dataset, [n_train, n_test], generator=torch.Generator().manual_seed(42))
37
38    trainloader = DataLoader(trainset, batch_size=batch_size)
39    testloader = DataLoader(testset, batch_size=batch_size)
40
41    return (trainloader, testloader, metric)
42
43
44class JKJ1Adataset(Dataset):
45    def __init__(self, 
46                 img_dir='data/JKJ1A-dataset/images', 
47                 label_dir='data/JKJ1A-dataset/labels',
48                 transform=None,
49                 target_transform=None,
50                 metric='beauty'):
51        
52        self.img_dir = img_dir
53        self.label_dir = os.path.join(label_dir, metric)
54
55        # 画像ファイルのパス一覧を取得する。
56        self.images = self.load_images(self.img_dir)
57        self.labels = np.load(os.path.join(self.label_dir, 'histogram.npy'))
58
59        self.transform = transform
60        self.target_transform = target_transform
61        # self.transform = transforms.Compose( 
62        #     [transforms.Resize(size), 
63        #      transforms.ToTensor(),
64        #      ] # original size is 256 x 256
65        #     )
66
67    def __getitem__(self, index):
68        img = self.images[index]
69        label = self.labels[index]
70        if self.transform is not None:
71            img = self.transform(img)
72        if self.target_transform is not None: 
73            label = self.target_transform(label)
74
75        return img, label
76
77    def _get_img_paths(self, img_dir, IMG_EXTENSIONS=[".jpg", ".jpeg", ".png", ".bmp"]):
78        """指定したディレクトリ内の画像ファイルのパス一覧を取得する。
79        """
80        img_dir = Path(img_dir)
81        img_paths = [
82            p for p in img_dir.iterdir() if p.suffix in IMG_EXTENSIONS
83        ]
84
85        return img_paths
86
87    def load_images(self, img_dir):
88        img_paths = sorted(self._get_img_paths(img_dir))
89        images = [Image.open(path) for path in img_paths]
90        return images
91
92    def _get_label_paths(self, label_dir):
93        """指定したディレクトリ内の画像ファイルのパス一覧を取得する。
94        """
95        label_dir = Path(label_dir)
96        label_paths = [
97            p for p in label_dir.iterdir()
98        ]
99
100        return label_paths
101
102    def __len__(self):
103        """ディレクトリ内の画像ファイルの数を返す。
104        """
105        return len(self.labels)
106