最大確率の低い点の選び方

1.前提・実現したいこと
Unetの出力の特徴マップ[4,64,256,256]からWとHの乱数を作成してランダムに1点選んでlossの計算に使っているのを、Unetの出力[4,4,256,256]をsoftmaxして最大確率が低い点から選んでlossの計算をするように変更したいです。

2.発生している問題
main.py loss.py Net.pyでフォルダを分けて作っているため、どこに書けばいいのか困っています。
3.該当のソースコード

Net.py
1# Unetの最後のブロックを抜粋
2# Block7  Decoder
3        h = F.relu(self.bnde3(self.deconv3(h)))  
4        h = torch.cat((h1, h), dim=1)                                 
5        h = F.relu(self.bn71(self.conv71(h)))
6        h = self.bn72(self.conv72(h))    #[4,64,256,256]          
7
8        # discriminate
9        y = self.conv_f(F.relu(h))       #[4,4,256,256]
10  
11        return y,h

loss.py
1import numpy as np
2import torch
3import torchvision
4import torch.nn as nn
5import torch.nn.functional as F
6from torchvision import datasets, transforms
7import os
8import argparse
9import random
10from tqdm import tqdm
11from sklearn.metrics import confusion_matrix
12from torch.autograd import Variable
13
14class SquareLoss(torch.nn.Module):
15    def __init__(self, device):
16        super(SquareLoss, self).__init__()
17        self.device = device
18
19    def forward(self, output, targets):  # output 出力(64次元)  targets ラベル
20        #output.size() => [b,c,W,H]
21        #targets.size() => [b,W,H]
22
23        #特徴マップF[c,W,H]から一つランダムに選ぶ -> 選んだものの大きさは[c]
24        ##ピクセル番号の乱数をつくる(2個) W,Hの2種類
25        
26        w = np.random.choice(256, 1)  # numpy配列をlist型に変換
27        w = w.tolist()
28        h = np.random.choice(256, 1)
29        h = h.tolist()
30        #print(output.shape)
31        #print(w,h)
32
33        #ベクトルのノルムを1にする
34        output = F.normalize(output, dim=1)
35        ##outputからピクセルの値を参照しデータを持ってくる 
36        a = output[:, :, w, h] 
37
38        #選んだものの大きさを[1,c,1,1]にして特徴マップFとcos類似度の計算
39        ##選んだものの大きさを[c] -> [1, c, 1, 1]にする
40        a = a.view(a.size(0), a.size(1), 1, 1)
41
42        ##以下は選んだ点と特徴マップFのcos類似度の計算

main.py
1from dataset import Covid19Loader
2from Loss_a import SquareLoss
3from Net import UNet
4省略
5######## IoUの測定関数 ########
6def IoU(output, targets, label):
7
8############## dataloader関数##############
9def dataload():
10
11############## train関数 ##############
12def train(epoch):
13    model.train()
14
15    sum_loss1 = 0   #提案手法のloss
16    sum_loss2 = 0   #softmax cross entropy loss
17    correct = 0
18    total = 0
19
20    # 学習ループ  inputs:入力画像  targets:教師ラベル
21    for batch_idx, (inputs, targets) in enumerate(tqdm(train_loader, leave=False)):
22 
23        inputs = inputs.cuda(device)
24        targets = targets.cuda(device)
25
26        targets = targets.long()
27
28        y,output = model(inputs)
29
30        # 損失計算   
31        loss1 = criterion1(output, targets)
32        loss2 = criterion2(y, targets)
33        loss = loss1 + loss2
34
35        # 勾配を0に
36        optimizer.zero_grad()
37
38        # 誤差逆伝搬 
39        loss.backward()
40
41        # パラメーター更新
42        optimizer.step()
43
44        # loss溜め
45        sum_loss1 += loss1.item()
46        sum_loss2 += loss2.item()
47
48        ###精度の計算###
49        # 出力をsoftmax関数に(0~1)
50        output = F.softmax(y, dim=1)
51
52        # 最大ベクトル
53        _, predicted = output.max(1)
54
55        # total = 正解, correct = 予測 
56        total += (targets.size(0)*targets.size(1)*targets.size(2))
57        correct += predicted.eq(targets).sum().item()  #predicted.eq(targets) : 教師ラベルと一致していたらTrue, 不一致ならFalse
58
59    return sum_loss1/(batch_idx+1), sum_loss2/(batch_idx+1), correct/total
60
61
62
63############## validation関数 ##############
64def val(epoch):
65#train関数とほとんど同じ