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

# + id="wn5ElkzkPMtd" colab_type="code" colab={}
# %load_ext autoreload
# %autoreload 2

# %matplotlib inline

# + [markdown] id="VbAfx7wYPMtk" colab_type="text"
# ## Does nn.Conv2d init work well?

# + id="UjbNOoguPMtn" colab_type="code" colab={}
#export
from exp.nb_02 import *

def get_data():
    path = Path('data/mnist.pkl.gz')
    with gzip.open(path, 'rb') as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')
    return map(tensor, (x_train,y_train,x_valid,y_valid))

def normalize(x, m, s): 
    return (x-m)/s


# + id="JVp1wq0CPMtq" colab_type="code" colab={}
# torch.nn.modules.conv._ConvNd.reset_parameters??

# + id="Iw78pjjfPMtt" colab_type="code" colab={}
x_train,y_train,x_valid,y_valid = get_data()
train_mean,train_std = x_train.mean(),x_train.std()
x_train = normalize(x_train, train_mean, train_std)
x_valid = normalize(x_valid, train_mean, train_std)

# + id="b0bfYueXPMtx" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="ab65985b-2a4a-4c54-8ae4-2fcb4956fdfd"
# resizing vector to 2d tensor to fit into convalution layer
x_train = x_train.view(-1,1,28,28)
x_valid = x_valid.view(-1,1,28,28)
x_train.shape,x_valid.shape

# + id="5zOeQB9QPMt2" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="2b98b5a6-5770-4c3c-e922-849141b83afc"
n,*_ = x_train.shape
c = y_train.max()+1
nh = 32
n,c

# + id="G_bxzJNMPMt5" colab_type="code" colab={}
l1 = nn.Conv2d(1, nh, 5) # 1 chanel input, 32 channel output with 5x5 kernel size

# + id="Hdzxtw06PMt7" colab_type="code" colab={}
# small sammple for testing purposes
x = x_valid[:100]

# + id="yIorwoINPMt-" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="d70cdce1-e759-4af8-af43-40b9a0ba298f"
x.shape


# + id="sI_FgoJBPMuE" colab_type="code" colab={}
def stats(x): 
    return x.mean(),x.std()


# + id="7Mt_SuqJPMuI" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="3b09d6e9-c04d-4469-a430-4520644d61ed"
l1.weight.shape

# + id="q9HsjCODPMuN" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 90} outputId="041c159d-da8f-49c6-c17d-49e5ec4d0395"
stats(l1.weight),stats(l1.bias)

# + id="nmJ0mzIaPMuR" colab_type="code" colab={}
t = l1(x)

# + id="az3Oik-yPMuV" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="6376a820-02d9-46cc-d431-a03e4883aa22"
stats(t)

# + [markdown] id="bdusftZQTB9Z" colab_type="text"
# standard deviation is not close to 1, so we compare the kaiming uniform with kaiming normal. 
#
# **Kaiming Normal uses Leaky ReLu layer**

# + id="phBY5jQsPMuZ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="f63be2c9-b4c0-4eff-9a2d-b9f8614aa590"
init.kaiming_normal_(l1.weight, a=1.)
stats(l1(x))

# + id="2noiiI1bPMuc" colab_type="code" colab={}
import torch.nn.functional as F


# + id="knBL_YMfPMuf" colab_type="code" colab={}
def f1(x,a=0): 
    return F.leaky_relu(l1(x),a)


# + id="JWSbKMO-PMui" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="003ad8cb-66e2-4ad7-a4df-5c7dcf61d135"
init.kaiming_normal_(l1.weight, a=0)
stats(f1(x))

# + id="QuFg_BUWPMuk" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="892cf47f-45ba-404b-ba70-30ac425279ff"
l1 = nn.Conv2d(1, nh, 5)
stats(f1(x))
# after relu it's even worse

# + id="XUehhaPIPMun" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="cf2cadf1-5def-4b14-a366-1454f19e58c1"
l1.weight.shape

# + [markdown] id="OMLlMey6Uuii" colab_type="text"
# **Implementing our own Kaiming initializer**

# + id="144DBAOkPMus" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="d3838ae8-9c54-4a0c-b9ff-5e8bd44a29bf"
# receptive field size; take one of the lement and call numel() will give the 5x5 num of elements.
rec_fs = l1.weight[0,0].numel()
rec_fs

# + id="a5MKnhCvPMux" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="0fd17713-15ac-49cf-9338-fd0a3d7149a6"
nf,ni,*_ = l1.weight.shape
nf,ni

# + id="YGToxrHIPMu1" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="c7e48d3e-0898-40df-c57f-0c7d15910c71"
fan_in  = ni*rec_fs # input channels * num of receptors
fan_out = nf*rec_fs # output channels * num of receptors
fan_in,fan_out


# + id="ARmlbjK-PMu4" colab_type="code" colab={}
# kaiming init formula
def gain(a): 
    return math.sqrt(2.0 / (1 + a**2))


# + id="LgImdqjaPMu7" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 108} outputId="0e9eed6b-7ec8-4108-b88f-3864a150639f"
gain(1),gain(0),gain(0.01),gain(0.1),gain(math.sqrt(5.))

# + id="oA4TbY0kPMu_" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="d6c0abef-a5d2-4c28-d1ae-8a5fc70991ad"
# we have implemented kaiming normal but cnv2d uses kaiming uniform; chacking the standard deviation for random numbers 
torch.zeros(10000).uniform_(-1,1).std()

# + id="6p_MNnWIPMvC" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="56c07709-2b6e-4dfe-f669-fad2ba99a8cf"
# equal to 1/√3; which we use to make normal distribution to uniform distribution
1/math.sqrt(3.)


# + id="hwZVEiYsPMvF" colab_type="code" colab={}
def kaiming2(x,a, use_fan_out=False):
    nf,ni,*_ = x.shape
    rec_fs = x[0,0].shape.numel()
    fan = nf*rec_fs if use_fan_out else ni*rec_fs
    std = gain(a) / math.sqrt(fan)
    bound = math.sqrt(3.) * std
    x.data.uniform_(-bound,bound)


# + id="EExh1l_BPMvN" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="6740cd51-0df5-4890-8b67-a939a09e39b9"
kaiming2(l1.weight, a=0);
stats(f1(x))

# + id="GAhjFxb3PMvq" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="16862067-acb0-4936-dfa5-b718bca73d77"
kaiming2(l1.weight, a=math.sqrt(5.)) # √5 used in pytorch
stats(f1(x))


# + id="pFeRO3H2PMvw" colab_type="code" colab={}
class Flatten(nn.Module):
    def forward(self,x): return x.view(-1)


# + id="uq3Lac4KPMv1" colab_type="code" colab={}
m = nn.Sequential(
    nn.Conv2d(1,8, 5,stride=2,padding=2), nn.ReLU(),
    nn.Conv2d(8,16,3,stride=2,padding=1), nn.ReLU(),
    nn.Conv2d(16,32,3,stride=2,padding=1), nn.ReLU(),
    nn.Conv2d(32,1,3,stride=2,padding=1),
    nn.AdaptiveAvgPool2d(1),
    Flatten(),
)

# + id="Ruhi6BkZPMv4" colab_type="code" colab={}
# sample dependent variables to test
y = y_valid[:100].float()

# + id="DMhK1LsYPMv9" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="1dea0ed0-41aa-44ff-ffcd-929bc5e79824"
t = m(x)
stats(t)

# + id="TPDTiS77PMv_" colab_type="code" colab={}
l = mse(t,y)
l.backward()

# + id="_OAoMl0NPMwC" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="f0f51fb8-4bc9-4c85-e062-6ce1ad532418"
stats(m[0].weight.grad)

# + id="j_DYh62gPMwG" colab_type="code" colab={}
# init.kaiming_uniform_??

# + id="s3GK5zwgPMwK" colab_type="code" colab={}
for l in m:
    if isinstance(l,nn.Conv2d):
        init.kaiming_uniform_(l.weight)
        l.bias.data.zero_()

# + id="GwtevPutPMwN" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 53} outputId="f7d2d8b2-3b36-4ecd-ced9-fd3f052e0607"
t = m(x)
stats(t)

# + id="Ml6xLM4-PMwQ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="a4eea6af-6d16-4193-f8d7-462eabfdac7a"
l = mse(t,y)
l.backward()
stats(m[0].weight.grad)

# + [markdown] id="vg1RQgfvPMwT" colab_type="text"
# ## Export

# + id="WFJySkw9PMwU" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="d2b2f647-c8c8-4909-a352-ce311a56c718"
# !chmod 755 ./notebook2script.py
# !./notebook2script.py 02a_why_sqrt5.ipynb

# + id="etit9VMHPMwY" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered