pytorch implementation of a model

#segmentation model on keras
puts = Input(shape=(256, 256, 1))

# encoding
net = Conv2D(32, kernel_size=3, activation='relu', padding='same')(inputs)
net = MaxPooling2D(pool_size=2, padding='same')(net)
net = Conv2D(64, kernel_size=3, activation='relu', padding='same')(net)
net = MaxPooling2D(pool_size=2, padding='same')(net)
net = Conv2D(128, kernel_size=3, activation='relu', padding='same')(net)
net = MaxPooling2D(pool_size=2, padding='same')(net)

net = Dense(128, activation='relu')(net)

# decoding
net = UpSampling2D(size=2)(net)
net = Conv2D(128, kernel_size=3, activation='sigmoid', padding='same')(net)
net = UpSampling2D(size=2)(net)
net = Conv2D(64, kernel_size=3, activation='sigmoid', padding='same')(net)
net = UpSampling2D(size=2)(net)

# output with 1 channel for gray scale segmenation
outputs = Conv2D(1, kernel_size=3, activation='sigmoid', padding='same')(net)

model = Model(inputs=inputs, outputs=outputs)

# use binary cross entropy with sigmoid function
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['acc', 'mse'])

model.summary()



#############################################################################################################################################################################################################################################################################################################################################################################################################

#model implementation on pytorch

import os
import torch
import torchvision
import tarfile
from torchvision.datasets.utils import download_url
from torch.utils.data import random_split
from torchsummary import summary



# Implementation of CNN/ConvNet Model

class build_unet(torch.nn.Module):

    def __init__(self):
        super(build_unet, self).__init__()
        # L1 ImgIn shape=(?, 28, 28, 1)
        # Conv -> (?, 28, 28, 32)
        # Pool -> (?, 14, 14, 32)
        keep_prob = 0.5
        self.layer1 = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, padding=1))
        # L2 ImgIn shape=(?, 14, 14, 32)
        # Conv      ->(?, 14, 14, 64)
        # Pool      ->(?, 7, 7, 64)
        self.layer2 = torch.nn.Sequential(
            torch.nn.Conv2d(32, 64, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, padding=1))
        # L3 ImgIn shape=(?, 7, 7, 64)
        # Conv ->(?, 7, 7, 128)
        # Pool ->(?, 4, 4, 128)
        self.layer3 = torch.nn.Sequential(
            torch.nn.Conv2d(64, 128, kernel_size=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, padding=1))



        self.dense = torch.nn.Linear(128, 128, bias=True)
        torch.nn.init.xavier_uniform_(self.dense.weight)
        self.layer4 = torch.nn.Sequential(
            self.dense,
            torch.nn.ReLU(),
            torch.nn.Upsample()
        )

        self.layer5 = torch.nn.Sequential(
            torch.nn.Conv2d(128, 128, kernel_size=3),
            torch.nn.Sigmoid(),
            torch.nn.Upsample()
        )

        self.layer6 = torch.nn.Sequential(
            torch.nn.Conv2d(128, 64, kernel_size=3),
            torch.nn.Sigmoid(),
            torch.nn.Upsample()
        )

        self.layer7 = torch.nn.Sequential(
            torch.nn.Conv2d(64, 1, kernel_size=3),
            torch.nn.Sigmoid()
        )



    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = self.layer6(out)
        out = self.layer7(out)


        #         out = out.view(out.size(0), -1)   # Flatten them for FC
#         out = self.fc1(out)
#         out = self.fc2(out)

        return out

if __name__ == "__main__":
    x = torch.randn((2, 3, 512, 512))
    f = build_unet()
    y = f(x)
    print(y.shape)