IndexError: index 256 is out of bounds for dimension 2 with size 256の対処方法

1.前提・実現したいこと
y[4,4,256,256]のsoftmaxから、最大確率が低いピクセルを選ぶようなプログラムの実装をしています。選んだピクセルはoutput[4,64,256,256]の特徴マップと演算をします。

2.発生している問題
学習の途中で以下のようなエラーが出て学習が止まってしまいました。
a = output[:, :, h, w]で
IndexError: index 256 is out of bounds for dimension 2 with size 256

3.該当のソースコード

#coding: utf-8
##### ライブラリ読み込み #####
import numpy as np
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
import os
import argparse
import random
from tqdm import tqdm
from sklearn.metrics import confusion_matrix
from torch.autograd import Variable

class SquareLoss(torch.nn.Module):
    def __init__(self, device):
        super(SquareLoss, self).__init__()
        self.device = device

    def forward(self, output, targets, y):   # output 出力(64次元)  targets ラベル
        #output.size() => [b,c,W,H]
        #targets.size() => [b,W,H]

        #初期設定
        h = 0
        w = 0
        
        output = F.normalize(output, dim=1)
        
        #一次元配列にする (b,c,H,W) -> (b,c,H*W) [4,4,65536]
        y = y.view(y.size(0), y.size(1), y.size(2) * y.size(3))  
        
        #softmaxをかけて確率にする
        ##256×256の一次元化されて列方向にsoftmaxされる
        y = F.softmax(y, dim=0)  #大きさ y.shape=[4, 4, 65536]
    
        #列方向の最大値を返す y1に値、y2にインデックスを取得する 65536がbatch個の行列
        y1,y2 = torch.max(y, dim=0, keepdim=False, out=None)
        #print(y1.shape) #[4, 65536]
        #print(y2.shape) #[4, 65536]

        #列方向の最小値を返す y1に値、y2にインデックスを取得する 65536が1個の行列
        y1,y2 = torch.min(y1, dim=0, keepdim=False, out=None)

        #最も確率が低い値のインデックスを取得する
        y3 = torch.argmin(y1, dim=None, keepdim=False)
        #print(y3)  #55772

        #y3を二次元の座標に変換する
        if y3 % output.size(2) == 0:
            h = y3 // output.size(2)
            w = 0
        else:
            h = y3 // output.size(2) + 1
            w = y3 % output.size(2)
        
        #print(h) #例 217
        #print(w)  #例 220
        
        #output(64次元の特徴マップ)からピクセルの値を参照しデータを持ってくる
        a = output[:, :, h, w]
        #print(a.shape) #ここでaの大きさはchannel [4,64]
        
        #aの大きさを[c] -> [1, c, 1, 1]にする
        a = a.view(a.size(0), a.size(1), 1, 1)
        #print(a.shape) [4,64,1,1]
    

        ##選んだ点と特徴マップを演算

4.自分で調べたこと
配列のサイズを超えてしまっているらしいのですがどこを修正したらよいのかわからないです。

5.補足情報
python3

otn

2021/09/22 04:32

> どこを修正したらよいのかわからないです。自分で内容を理解していないプログラムだと言うことでしょうか？

mek_41

2021/09/22 08:03

二次元に変換するときにインデックスが0からカウントされることを考慮してなかったことで出たエラーでした。修正し直しができたので良かったです。触り始めたばかりで自分で書いたコードでも完璧に理解しきれてないところはあると思います。teratailはわからないところは的確にアドバイス・回答してくださる方が多く感謝しています。今後は質問の仕方も具体的に書けるように気を付けたいと思います、コメントしていただきありがとうございました!

otn

2021/09/22 09:13

自力で解決できたようで良かったです。