VAEを作ったときのサイズを合わせるためのツールスクリプト

vae_check.py

import numpy as np

class Sequential:
    def __init__(self, channel=0, height=0, width=0, layers=[], source=True):
        print('--------------------------------')
        if source:
            for layer in layers:
                (channel, height, width) = layer.source(channel, height, width)
        else:
            for layer in layers:
                (channel, height, width) = layer.size(channel, height, width)

        print((channel, height, width))

class Conv2d:
    def __init__(self, channel_input, channel_output, stride, kernel, padding, dialation=(1,1)):
        self.channel_input = channel_input
        self.channel_output = channel_output
        self.stride = stride
        self.kernel = kernel
        self.padding = padding
        self.dialation = dialation
        self.name = "Conv2d          "

    def through(self, height, width):
        temp = (height + 2 * self.padding[0] - self.dialation[0] * (self.kernel[0] - 1) - 1)
        height_output = int(temp / self.stride[0]) + 1
        temp = (width  + 2 * self.padding[1] - self.dialation[1] * (self.kernel[1] - 1) - 1)
        width_output = int(temp / self.stride[1]) + 1
        return (height_output, width_output)

    def source(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        value = (channel_input, self.channel_output, self.padding[0], self.padding[1], self.kernel[0], self.kernel[1], self.stride[0], self.stride[1])
        print("nn.Conv2d(%3d, %d, padding=(%d, %d), kernel_size=(%d, %d), stride=(%d, %d))," % value)
        return (self.channel_output, height_output, width_output)
    
    def size(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        print("%s%3dx%3dx%3d -> %3dx%3dx%3d" % (self.name, channel_input, height_input, width_input, self.channel_output, height_output, width_output))
        return (self.channel_output, height_output, width_output)

class ConvTranspose2d:
    def __init__(self, channel_input, channel_output, stride, kernel, padding, dialation=(1,1)):
        self.channel_input = channel_input
        self.channel_output = channel_output
        self.stride = stride
        self.kernel = kernel
        self.padding = padding
        self.dialation = dialation
        self.name = "ConvTranspose2d "

    def through(self, height, width):
        height_output = (height - 1) * self.stride[0] - 2 * self.padding[0] + self.dialation[0] * (self.kernel[0] - 1) + self.padding[0] + 1
        width_output  = (width  - 1) * self.stride[1] - 2 * self.padding[1] + self.dialation[1] * (self.kernel[1] - 1) + self.padding[1] + 1
        return (height_output, width_output)

    def source(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        value = (channel_input, self.channel_output, self.padding[0], self.padding[1], self.kernel[0], self.kernel[1], self.stride[0], self.stride[1])
        print("nn.ConvTranspose2d(%3d, %3d, padding=(%d, %d), kernel_size=(%d, %d), stride=(%d, %d))," % value)
        return (self.channel_output, height_output, width_output)

    def size(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        print("%s%3dx%3dx%3d -> %3dx%3dx%3d" % (self.name, channel_input, height_input, width_input, self.channel_output, height_output, width_output))
        return (self.channel_output, height_output, width_output)

class Upsample:
    def __init__(self, scale):
        self.scale = scale
        self.name = "Upsample        "

    def through(self, channel, height, width):
        channel_output = channel
        height_output = height * self.scale
        width_output =  width * self.scale
        return (channel_output, height_output, width_output)

    def source(self, channel_input, height_input, width_input):
        (channel_output, height_output, width_output) = self.through(channel_input, height_input, width_input)
        print("nn.Upsample(scale_factor=%d, mode='bilinear', align_corners=True)," % self.scale)
        return (channel_output, height_output, width_output)

    def size(self, channel_input, height_input, width_input):
        (channel_output, height_output, width_output) = self.through(channel_input, height_input, width_input)
        print("%s%3dx%3dx%3d -> %3dx%3dx%3d" % (self.name, channel_input, height_input, width_input, channel_output, height_output, width_output))
        return (channel_output, height_output, width_output)

class ReLU:
    def __init__(self):
        self.name = "ReLU            "

    def through(self, height, width):
        height_output = height
        width_output =  width
        return (height_output, width_output)

    def source(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        print("nn.ReLU(),")
        return (channel_input, height_output, width_output)

    def size(self, channel_input, height_input, width_input):
        (height_output, width_output) = self.through(height_input, width_input)
        print("%s%3dx%3dx%3d -> %3dx%3dx%3d" % (self.name, channel_input, height_input, width_input, channel_input, height_output, width_output))
        return (channel_input, height_output, width_output)

for source_flag in [False, True]:
    Sequential(
        channel=3,
        height=128,
        width=128,
        source=source_flag,
        layers=[
            Conv2d(  3, 32, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
            Conv2d( 32, 64, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
            Conv2d( 64,128, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
            Conv2d(128,128, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
            Conv2d(128,128, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
            Conv2d(128,128, stride=(2, 2), kernel=(4, 4), padding=(1, 1)),   ReLU(),
        ]
    )

    Sequential(
        channel=512,
        height=1,
        width=1,
        source=source_flag,
        layers=[
            ConvTranspose2d(512, 256, stride=(1, 1), kernel=(4, 4), padding=(0, 0)), ReLU(),
            Upsample(2), ReLU(),
            Conv2d(256,256, stride=(1, 1), kernel=(3, 3), padding=(1, 1)),  ReLU(),
            Upsample(2), ReLU(),
            Conv2d(256,256, stride=(1, 1), kernel=(3, 3), padding=(1, 1)),  ReLU(),
            Upsample(2), ReLU(),
            Conv2d(256,256, stride=(1, 1), kernel=(3, 3), padding=(1, 1)),  ReLU(),
            Upsample(2), ReLU(),
            Conv2d(256,256, stride=(1, 1), kernel=(3, 3), padding=(1, 1)),  ReLU(),
            Upsample(2), ReLU(),
            Conv2d(256,3, stride=(1, 1), kernel=(3, 3), padding=(1, 1)),  ReLU(),
            
        ]
    )