Convolution1Dの入出力サイズを指定して、必要なカーネルサイズとストライドを計算するコード

convolution層を設計する時に、jupyter上でちまちまパラメータを変えてshapeを何度も確認するという事を繰り返していたのですが、
こんなのは人間のやる仕事ではない!何とか楽にしよう。という事で、タイトルの内容のコードを書く事にしました。
(でも、もっと上手い方法あるかも)
私は、VAEの設計をしているのでEncoderとDecoderの例を載せておきます。

# Encoder設計ツール
import itertools

# kernelとstrideの組み合わせを全て試す
# 範囲を広くとるとめちゃくちゃ出力されるので注意!
k = range(8, 10)
s = range(1, 3)

input_size = 600
target_size = 140

for k1,k2,k3,k4,s1,s2,s3,s4 in itertools.product(k, k, k, k, s, s, s, s):

    Conv1 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=k1, stride=s1, padding=0)
    Conv2 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=k2, stride=s2, padding=0)
    Conv3 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=k3, stride=s3, padding=0)
    Conv4 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=k4, stride=s4, padding=0)

    x = torch.randn(1, 1, input_size)
    x = Conv1(x)
    x = Conv2(x)
    x = Conv3(x)
    if k4 < x.squeeze().shape[0]:
        x = Conv4(x)
    x = x.squeeze().shape[0]

    if x == target_size:
        print("k1 : ",k1,"k2 : ",k2,"k3 : ",k3,"k4 : ",k4,"s1 : ",s1,"s2 : ",s2,"s3 : ",s3,"s4 : ",s4)

print("finished")  
# Decoder設計ツール
import itertools

# kernelとstrideの組み合わせを全て試す
# 範囲を広くとるとめちゃくちゃ出力されるので注意!
k = range(4, 6)
s = range(1, 3)

input_size = 140
target_size = 600

for k1,k2,k3,k4,s1,s2,s3,s4 in itertools.product(k, k, k, k, s, s, s, s):

    Conv1 = nn.ConvTranspose1d(in_channels=1, out_channels=1, kernel_size=k1, stride=s1, padding=0)
    Conv2 = nn.ConvTranspose1d(in_channels=1, out_channels=1, kernel_size=k2, stride=s2, padding=0)
    Conv3 = nn.ConvTranspose1d(in_channels=1, out_channels=1, kernel_size=k3, stride=s3, padding=0)
    Conv4 = nn.ConvTranspose1d(in_channels=1, out_channels=1, kernel_size=k4, stride=s4, padding=0)

    x = torch.randn(1, 1, input_size)
    x = Conv1(x)
    x = Conv2(x)
    x = Conv3(x)
    if k4 < x.squeeze().shape[0]:
        x = Conv4(x)
    x = x.squeeze().shape[0]
    #print(x)
    if x == target_size:
        print("k1 : ",k1,"k2 : ",k2,"k3 : ",k3,"k4 : ",k4,"s1 : ",s1,"s2 : ",s2,"s3 : ",s3,"s4 : ",s4)

print("finished")  

使い方は、コード見ればわかると思うので割愛
別に改造すれば、Conv2Dにも使えるしパディングサイズを可変出来る様になるので、私にはそのモチベーションはないけれど誰か必要な人のベースになればいいなと
限定的な使い方しかできないけど、個人的なストレスはかなり軽減されたので載せておきます。パラパラっと書いたものなので至らない点があるかもですがご容赦を

Google Colabのリンク
https://colab.research.google.com/drive/1arImaQrMqmZwhQi350bCt3h4MGYLd-mL?usp=sharing


いいなと思ったら応援しよう!