[pytorchモデル]forwardの中身が知りたい。特にsoftmaxのdimの決め方がわからない。

やりたいこと

以下の条件でモデルを作成したい

・pytorchであること
・画像データとlabelを付与することについてやっています。

・入力画像は256x256のRGB画像
・出力は二値で分類するので2次元になると思います

環境

python 3.7.7

>>> print(torch.__version__)
1.5.1+cu101

書いてみたコード



import torch
import torch.nn as nn
import torch.nn.functional as F

"""
入力は256x256のRGB画像
出力は0or1の2次元ベクトルになるはず

"""



class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()

        # input: RGB 3 ,output : 16, kernel 3(default)
        self.conv1 = nn.Conv2d(3, 16, 3)
        # print(self.conv1.weight)
        self.pool1 = nn.MaxPool2d(2,stride=2)

        self.conv2 = nn.Conv2d(16, 32, 3)
        self.pool2 = nn.MaxPool2d(2,stride=2)

        self.conv3 = nn.Conv2d(32, 64, 3)
        self.pool3 = nn.MaxPool2d(2,stride=2)

        self.conv4 = nn.Conv2d(64, 128, 3)
        self.pool4 = nn.MaxPool2d(2,stride=2)

        self.conv5 = nn.Conv2d(128, 256, 3)
        self.pool5 = nn.MaxPool2d(2,stride=2)

        self.softmax = nn.Softmax(dim=1)



        # an affine operation: y = Wx + b
        # linearにおけるout_featuresとはなんだろうか
        self.fc1 = nn.Linear(256 * 3 * 3, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2)

    def forward(self, x):
        # conv relu poolの流れを実行してみる

        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool2(x)
       
        x = self.conv3(x)
        x = self.relu(x)
        x = self.pool3(x)

        x = self.conv4(x)
        x = self.relu(x)
        x = self.pool4(x)

        x = self.conv5(x)
        x = self.relu(x)
        x = self.pool5(x)

        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        x = self.softmax(x)

        return x




net = Net()

print(net)
# params = list(net.parameters())
# print(len(params))

出力

Net(
  (relu): ReLU()
  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1))
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv2): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1))
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv3): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (pool3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv4): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
  (pool4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (conv5): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1))
  (pool5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (softmax): Softmax(dim=1)
  (fc1): Linear(in_features=2304, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=2, bias=True)
)

＜コードの意図＞
conv→relu→poolingのユニットを自分にとってわかりやすくするためにまとめて書いています。
convの引数の意図は最初のinはRGBなので3、kernelは３で統一しています。最終的に最初の256になるように大きくしていきます。

また、Linearの意図に関しては256から2次元にするために最初と最後は決まっていますが、
間にどういった法則性で小さくしていくかわからなかったため120、84のような適当な値になってしまいました。
softmaxの書き方がどうなっているのかわからなかったのでdim=1にしています。