IndexError: Target 534 is out of bounds.

前提・実現したいこと

自作データセット(1006のクラス、Train＝20枚、Test=4枚、32×32のpng画像)を用いて、文字認識をするモデルを作成しようとしていた時に、タイトルのようなエラーがでてしまいました。

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

IndexError                                Traceback (most recent call last)
<ipython-input-20-02d64264044b> in <module>()
     10         optimizer.zero_grad()
     11         outputs = model(inputs)
---> 12         loss = criterion(outputs, labels)
     13         loss.backward()
     14         optimizer.step()

3 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/loss.py in forward(self, input, target)
    960     def forward(self, input: Tensor, target: Tensor) -> Tensor:
    961         return F.cross_entropy(input, target, weight=self.weight,
--> 962                                ignore_index=self.ignore_index, reduction=self.reduction)
    963 
    964 

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
   2466     if size_average is not None or reduce is not None:
   2467         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2468     return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
   2469 
   2470 

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2262                          .format(input.size(0), target.size(0)))
   2263     if dim == 2:
-> 2264         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2265     elif dim == 4:
   2266         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
IndexError: Target 534 is out of bounds.

該当のソースコード

python
1class MyNet(nn.Module):
2    def __init__(self):
3        super(MyNet, self).__init__()
4        self.relu = nn.ReLU()
5        self.softmax = nn.Softmax(1)
6
7        self.conv1 = nn.Conv2d(1,16,5)
8        self.conv2 = nn.Conv2d(16,32,5)
9        self.conv3 = nn.Conv2d(32,64,5)
10
11        self.pool = nn.MaxPool2d(2)
12
13        self.flatten = nn.Flatten()
14
15        self.fc1 = nn.Linear(64 * 1 * 1, 100)
16        self.fc2 = nn.Linear(100,10)
17
18    def forward(self, x):
19      x = self.conv1(x)
20      x = self.relu(x)
21      x = self.pool(x)
22      x = self.conv2(x)
23      x = self.relu(x)
24      x = self.pool(x)
25      x = self.conv3(x)
26      x = self.relu(x)
27      x = self.flatten(x)
28      x = self.fc1(x)
29      x = self.relu(x)
30      x = self.fc2(x)
31      #x = self.softmax(x)
32      return x
33
34args = parse()
35phase = 'train' or 'test'
36mean = (0.5,)
37std  = (0.5,)
38transform_dict = {
39        'train': transforms.Compose(
40            [transforms.Grayscale(),
41             transforms.Resize((32,32)),
42             #transforms.RandomHorizontalFlip(),
43             transforms.ToTensor(),
44             transforms.Normalize(mean,std)
45             #transforms.Normalize(mean=[0.485, 0.456, 0.406],
46                                  #std=[0.229, 0.224, 0.225]),
47             ]),
48        'test': transforms.Compose(
49            [transforms.Grayscale(),
50             transforms.Resize((32,32)),
51             transforms.ToTensor(),
52             transforms.Normalize(mean,std)
53            ])}
54
55def prepare_dataset(transform, transform_val, dirname_dataset, dirname_dataset_val):
56  dataset = torchvision.datasets.ImageFolder(dirname_dataset, transform)
57  dataset_val = torchvision.datasets.ImageFolder(dirname_dataset_val, transform_val)
58  print(dataset.class_to_idx)
59  print(dataset_val.class_to_idx)
60  return dataset, dataset_val
61
62train_data, test_data = prepare_dataset(transform_dict["train"], transform_dict["test"], "/content/drive/MyDrive/Train", "/content/drive/MyDrive/Test")
63
64train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batchsize, shuffle=True)
65val_loader   = torch.utils.data.DataLoader(test_data, batch_size=args.batchsize, shuffle=False)
66dataloaders  = {"train":train_loader, "val":val_loader}
67
68model = MyNet()
69criterion = nn.CrossEntropyLoss()
70optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9, weight_decay=0.005)
71
72epochs = 2
73running_loss = 0.0
74
75for epoch in range(epochs):
76    for (inputs, labels) in train_loader:
77
78        # GPUを使用しない場合は削除
79        #inputs, labels = inputs.to(device), labels.to(device)
80
81        optimizer.zero_grad()
82        outputs = model(inputs)
83        loss = criterion(outputs, labels)
84        loss.backward()
85        optimizer.step()
86
87        running_loss += loss.item()
88
89    print("epoch {}/{}, loss={}".format(epoch + 1, epochs, running_loss / 2000))