# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Implementation of Resnet 34.

from importlib.util import find_spec
if find_spec("model") is None:
    import sys
    sys.path.append('..')

# ### Pytorch Implementation.

import torch.nn as nn
import torch.nn.functional as F
from base import BaseModel


class Resnet34(BaseModel):
    def __init__(self, num_classes=100):
        super().__init__()
        
        ## num_flatten_params = 512 * 1 * 1
        num_channels = 512
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=2)
        self.pool1 = nn.MaxPool2d(stride=2, kernel_size=2)
        
        
        self.res_block1 = self._build_residual_block(64, 64)
        
        self.res_block2 = self._build_residual_block(64, 64)
        
        self.res_block3 = self._build_residual_block(64, 64)
        
        self.res_block4 = self._build_residual_block(64, 128, stride=2)
        
        self.res_block5 = self._build_residual_block(128, 128)
        
        self.res_block6 = self._build_residual_block(128, 128)
        
        self.res_block7 = self._build_residual_block(128, 128)
        
        self.res_block8 = self._build_residual_block(128, 256, stride=2)
        
        self.res_block9 = self._build_residual_block(256, 256)
        
        self.res_block10 = self._build_residual_block(256, 256)
        
        self.res_block11 = self._build_residual_block(256, 256)
        
        self.res_block12 = self._build_residual_block(256, 256)
        
        self.res_block13 = self._build_residual_block(256, 256)
        
        self.res_block14 = self._build_residual_block(256, 512, stride=2)
        
        self.res_block15 = self._build_residual_block(512, 512)
        
        self.res_block16 = self._build_residual_block(512, 512)
        
        self.global_avg_pooling = nn.AvgPool2d(kernel_size=7)
        
        self.fc = nn.Linear(num_channels, num_classes)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        assert(x.shape == (x.size(0), 112, 112, 64))
        
        x = self.pool1(x)
        assert(x.shape == (x.size(0), 56, 56, 64))
        
        x = self.res_block1(x) + x
        x = self.res_block2(x) + x
        x = self.res_block3(x) + x
        assert(x.shape == (x.size(0), 56, 56, 64))
        
        x = self.res_block4(x) + self._zero_pad(x, 64)
        assert(x.shape == (x.size(0), 28, 28, 128))
        x = self.res_block5(x) + x
        x = self.res_block6(x) + x
        x = self.res_block7(x) + x
        
        x = self.res_block8(x) + self._zero_pad(x, 128)
        assert(x.shape == (x.size(0), 14, 14, 256))
        x = self.res_block9(x) + x
        x = self.res_block10(x) + x
        x = self.res_block11(x) + x
        x = self.res_block12(x) + x
        x = self.res_block13(x) + x
        
        x = self.res_block14(x) + self._zero_pad(x, 256)
        assert(x.shape == (x.size(0), 7, 7, 512))
        x = self.res_block15(x) + x
        x = self.res_block16(x) + x
        
        x = self.global_avg_pooling(x)
        assert(x.shape == (x.size(0), 512))
        
        return x.fc(x)
        
    def _build_residual_block(self, in_channels, out_channels, stride=1, padding=1):
        return nn.Sequential(
            self._build_conv_layer(in_channels, out_channels, stride, padding),
            nn.ReLU(inplace=True),
            self._build_conv_layer(out_channels, out_channels),
            nn.ReLU(inplace=True)
        )
    
    def _zero_pad(self, x, amount):
        return F.pad(x, (amount, amount), "constant", 0)
    
    def _build_conv_layer(self, in_channels, out_channels, stride=1, padding=1):
        return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=padding)
    
    def _flatten(self, x):
        return x.view(x.size(0) - 1)
    

model = Resnet34()
model.parameters

# ## Keras Implementation

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers



StarcoderdataJupyterScriptsDedupFiltered