torch.nn.functional.normalizeの正しいパラメーターの設定の仕方

1.前提・実現したいこと
CNN(UNet)で抽出した特徴量のcos類似度の計算を行っています。
具体的にはoutput[4,64,256,256]から1つ選び、その点とoutputの全ての点とのcos類似度を計算します。
output[4,64,256,256]ベクトルのノルムを１にしたいです。

2.発生している問題
output[4,64,256,256]のベクトルのノルムを1にしたいのですが、パラメーターの設定の仕方が分からず困っています。

3.該当のソースコード

cos類似度の計算
1##### ライブラリ読み込み #####
2import numpy as np
3import torch
4import torchvision
5import torch.nn as nn
6import torch.nn.functional as F
7from torchvision import datasets, transforms
8import os
9import argparse
10import random
11from tqdm import tqdm
12from sklearn.metrics import confusion_matrix
13from torch.autograd import Variable
14############################
15
16class SquareLoss(torch.nn.Module):
17    def __init__(self, device):
18        super(SquareLoss, self).__init__()
19        self.device = device
20
21    def forward(self, output, targets):   # output 出力  targets ラベル
22        #output.size() => [b,c,W,H]
23        #targets.size() => [b,W,H]
24
25        #特徴マップF[b,c,W,H]から一つランダムに選ぶ -> 選んだものの大きさは[c]
26        ##ピクセル番号の乱数をつくる(2個) W,Hの2種類
27        
28        w = np.random.choice(256, 1)  # numpy配列をlist型に変換
29        w = w.tolist()
30        h = np.random.choice(256, 1)
31        h = h.tolist()
32        
33        ##outputからピクセルの値を参照しデータを持ってくる ピクセルの番号->値
34        a = output[:, :, w, h] 
35        #print(a.shape) ここでaの大きさはchannel
36
37        #選んだものの大きさを[1,c,1,1]にして特徴マップFと演算(cos類似度)
38        #->勝手に全ての点の計算をしてくれる
39        ##選んだものの大きさを[c] -> [1, c, 1, 1]にする
40        a = a.view(a.size(0), a.size(1), 1, 1)
41
42        ##選んだ点と特徴マップFを演算(cos類似度の計算)
43        ###分子の計算(内積計算のみ)
44        #print(a.shape)               # [4, 64, 1, 1]
45        #print(output.shape)          # [4, 64, 256, 256]
46        num = (a*output).sum(dim=1)
47
48        ###分母の計算
49        ###ベクトルaの絶対値を計算
50        a_sq = a ** 2
51        
52        ###要素が2乗されたベクトルaの要素の和
53        a_sum = a_sq.sum(dim=1)                
54        
55        ###上で得られた値の平方根をとり、aの絶対値を取得
56        a_root = torch.sqrt(torch.clamp(a_sum, min=1.0e-6))
57
58        ###特徴ベクトル(output)のノルムを1にする
59        **output = F.normalize(output, p=2, dim=1, eps=1e-12, out=None)**
60        
61        ###同様の操作をベクトルoutputにも行う
62        output_sq = output ** 2
63        output_sum = output_sq.sum(dim=1)
64        output_root = torch.sqrt(torch.clamp(output_sum, min=1.0e-6))
65
66        den = a_root * output_root
67        
68        ###分子を分母で割ってcos類似度を求める
69        cos_sim = num / den  # -1まで出力される
70        
71        ###出力
72        #print(cos_sim)
73

4.自分で調べたことや試したこと
torch.nn.functional.nornalize(入力、p = 2.0、dim = 1、eps = 1e-12、out =なし）を使うことで指定した入力の正規化ができることはわかりましたが、パラメーター(pやdim)の設定が分からないです。
python始めたばかりでわからないことが初歩的なことですいません。

5.補足情報
python3

toast-uz

2021/09/24 23:37

すみません。質問の意味が分かりません。正しいパラメーターの設定の仕方、という表題にもかかわらず、指定した入力の正規化ができた = 正しいパラメーターの設定の仕方がわかった、という内容になっていると思いますが、何を質問されているのでしょうか？パラメータの意味が知りたい、ということでしょうか？また、PyTorchのドキュメントでパラメータの定義が書かれていますが、どこがわからないのでしょうか？