witnn/CNNDemo/Resnet50.py

from __future__ import print_function

import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
import torchvision.models as models
import matplotlib.pyplot as plt
import numpy as np
import struct
from struct import Struct

CurrentPath = os.path.split(os.path.realpath(__file__))[0]+"/"

# resnet50 = models.resnet50(pretrained=True)
# torch.save(resnet50, CurrentPath+'params.pth')
resnet50 = torch.load(CurrentPath+'params.pth')
resnet50.eval()


# print("===========================")
# print("===========================")
# print("===========================")
# print(resnet50)


ResNet50 = {
    "conv1": "Conv2d",
    "bn1": "BatchNorm2d",
    "relu": "ReLU",
    "maxpool": "MaxPool2d",
    "layer1": {
        "_modules": {
            "0": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
                "downsample": {
                    "_modules": {
                        "0": "Conv2d",
                        "1": "BatchNorm2d",
                    }
                }
            },
            "1": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "2": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
        }
    },
    "layer2": {
        "_modules": {
            "0": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
                "downsample": {
                    "_modules": {
                        "0": "Conv2d",
                        "1": "BatchNorm2d",
                    }
                }
            },
            "1": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "2": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "3": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            }
        }
    },
    "layer3": {
        "_modules": {
            "0": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
                "downsample": {
                    "_modules": {
                        "0": "Conv2d",
                        "1": "BatchNorm2d",
                    }
                }
            },
            "1": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "2": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "3": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "4": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "5": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            }
        }
    },
    "layer4": {
        "_modules": {
            "0": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
                "downsample": {
                    "_modules": {
                        "0": "Conv2d",
                        "1": "BatchNorm2d",
                    }
                }
            },
            "1": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            },
            "2": {
                "conv1": "Conv2d",
                "bn1": "BatchNorm2d",
                "conv2": "Conv2d",
                "bn2": "BatchNorm2d",
                "conv3": "Conv2d",
                "bn3": "BatchNorm2d",
                "relu": "ReLU",
            }
        }
    },
    "avgpool": "AdaptiveAvgPool2d",
    "fc": "Linear"
}


weightfile = open(CurrentPath+'ResNet50Weight.cc', 'w')
shapefile = open(CurrentPath+'ResNet50WeightShape.cc', 'w')
binaryfile = open(CurrentPath+'ResNet50Weight.bin', 'wb')
currentbyte = 0
strg = ''
strgShape = ''

def genData(name, data):
    global currentbyte
    global binaryfile
    global strg
    global strgShape

    strg = strg + "int "+name+"[] = { "
    array = data.cpu().detach().numpy().reshape(-1)
    strg += str(currentbyte) + ","
    for a in array:
            bs = struct.pack("f", a)
            binaryfile.write(bs)
            currentbyte = currentbyte+4
    strg += str(currentbyte-1)
    strg = strg + " };\n"

    strgShape = strgShape + "int "+name+"_shape[] = { "
    array = data.cpu().detach().numpy().shape
    for a in array:
        strgShape += str(a) + ", "        
    strgShape = strgShape + " };\n"



def hook_fn(m, i, o):
  print(m)
  print("------------Input Grad------------")

  for grad in i:
    try:
      print(grad.shape)
    except AttributeError: 
      print ("None found for Gradient")

  print("------------Output Grad------------")
  for grad in o:  
    try:
      print(grad.shape)
    except AttributeError: 
      print ("None found for Gradient")
  print("\n")


def hook_print(name, m, i, o):
    global currentbyte
    global binaryfile
    global strg
    genData(name+"_input", i[0])
    genData(name+"_output", o[0])


def printDick(d, head, obj):
  global strg
  for item in d:
    if type(d[item]).__name__ == 'dict':
        objsub = getattr(obj, item, '')
        if objsub == '':
            objsub = obj[item]
        printDick(d[item], head+"_"+item, objsub)
    else:
        objsub = getattr(obj, item, '')
        if objsub == '':
            objsub = obj[item]
        if d[item] == "Conv2d":
            genData(head+"_"+item+"_weight", objsub.weight)
            strg = strg + "\n"

        if d[item] == "BatchNorm2d":
            genData(head+"_"+item+"_running_mean", objsub.running_mean)
            genData(head+"_"+item+"_running_var", objsub.running_var)
            genData(head+"_"+item+"_weight", objsub.weight)
            genData(head+"_"+item+"_bias", objsub.bias)
            strg = strg + "\n"

        if d[item] == "Linear":
            genData(head+"_"+item+"_weight", objsub.weight)
            genData(head+"_"+item+"_bias", objsub.bias)
            strg = strg + "\n"

printDick(ResNet50, "RN50", resnet50)

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

val_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(CurrentPath+'ImageNet/', transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        normalize,
    ])),
    batch_size=1, shuffle=False,
    num_workers=1, pin_memory=True)

strg = strg + "\n"
strg = strg + "\n"
strg = strg + "\n"
strg = strg + "\n"
strg = strg + "// val data 0-9 \n"

for batch_idx, (data, target) in enumerate(val_loader):
    genData("input_"+str(batch_idx), data)
    out = resnet50(data)
    genData("output_"+str(batch_idx), out)


strg = strg + "\n"
strg = strg + "\n"
strg = strg + "\n"
strg = strg + "\n"
strg = strg + "// input 0 layer output \n"

for batch_idx, (data, target) in enumerate(val_loader):
    genData("verify_input", data)
    x = resnet50.conv1(data)
    genData("verify_conv1", x)
    x = resnet50.bn1(x)
    genData("verify_bn1", x)
    x = resnet50.relu(x)
    genData("verify_relu", x)
    x = resnet50.maxpool(x)
    genData("verify_maxpool", x)
    x = resnet50.layer1(x)
    genData("verify_layer1", x)
    x = resnet50.layer2(x)
    genData("verify_layer2", x)
    x = resnet50.layer3(x)
    genData("verify_layer3", x)
    x = resnet50.layer4(x)
    genData("verify_layer4", x)
    x = resnet50.avgpool(x)
    genData("verify_avgpool", x)
    x = torch.flatten(x, 1)
    x = resnet50.fc(x)
    genData("verify_fc", x)
    break



resnet50.layer1._modules['0'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_bn1", m, i, o))
resnet50.layer1._modules['0'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_bn2", m, i, o))
resnet50.layer1._modules['0'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_bn3", m, i, o))
resnet50.layer1._modules['0'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_conv1", m, i, o))
resnet50.layer1._modules['0'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_conv2", m, i, o))
resnet50.layer1._modules['0'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer1_block0_conv3", m, i, o))

resnet50.layer1._modules['0'].downsample._modules['0'].register_forward_hook(lambda m, i, o: hook_print("layer1_block0_downsample_conv", m, i, o))
resnet50.layer1._modules['0'].downsample._modules['1'].register_forward_hook(lambda m, i, o: hook_print("layer1_block0_downsample_bn", m, i, o))

resnet50.layer1._modules['1'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_bn1", m, i, o))
resnet50.layer1._modules['1'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_bn2", m, i, o))
resnet50.layer1._modules['1'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_bn3", m, i, o))
resnet50.layer1._modules['1'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_conv1", m, i, o))
resnet50.layer1._modules['1'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_conv2", m, i, o))
resnet50.layer1._modules['1'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer1_block1_conv3", m, i, o))

resnet50.layer1._modules['2'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_bn1", m, i, o))
resnet50.layer1._modules['2'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_bn2", m, i, o))
resnet50.layer1._modules['2'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_bn3", m, i, o))
resnet50.layer1._modules['2'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_conv1", m, i, o))
resnet50.layer1._modules['2'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_conv2", m, i, o))
resnet50.layer1._modules['2'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer1_block2_conv3", m, i, o))




resnet50.layer2._modules['0'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_bn1", m, i, o))
resnet50.layer2._modules['0'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_bn2", m, i, o))
resnet50.layer2._modules['0'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_bn3", m, i, o))
resnet50.layer2._modules['0'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_conv1", m, i, o))
resnet50.layer2._modules['0'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_conv2", m, i, o))
resnet50.layer2._modules['0'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer2_block0_conv3", m, i, o))

resnet50.layer2._modules['0'].downsample._modules['0'].register_forward_hook(lambda m, i, o: hook_print("layer2_block0_downsample_conv", m, i, o))
resnet50.layer2._modules['0'].downsample._modules['1'].register_forward_hook(lambda m, i, o: hook_print("layer2_block0_downsample_bn", m, i, o))

resnet50.layer2._modules['1'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_bn1", m, i, o))
resnet50.layer2._modules['1'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_bn2", m, i, o))
resnet50.layer2._modules['1'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_bn3", m, i, o))
resnet50.layer2._modules['1'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_conv1", m, i, o))
resnet50.layer2._modules['1'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_conv2", m, i, o))
resnet50.layer2._modules['1'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer2_block1_conv3", m, i, o))

resnet50.layer2._modules['2'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_bn1", m, i, o))
resnet50.layer2._modules['2'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_bn2", m, i, o))
resnet50.layer2._modules['2'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_bn3", m, i, o))
resnet50.layer2._modules['2'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_conv1", m, i, o))
resnet50.layer2._modules['2'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_conv2", m, i, o))
resnet50.layer2._modules['2'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer2_block2_conv3", m, i, o))

resnet50.layer2._modules['3'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_bn1", m, i, o))
resnet50.layer2._modules['3'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_bn2", m, i, o))
resnet50.layer2._modules['3'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_bn3", m, i, o))
resnet50.layer2._modules['3'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_conv1", m, i, o))
resnet50.layer2._modules['3'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_conv2", m, i, o))
resnet50.layer2._modules['3'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer2_block3_conv3", m, i, o))




resnet50.layer3._modules['0'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_bn1", m, i, o))
resnet50.layer3._modules['0'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_bn2", m, i, o))
resnet50.layer3._modules['0'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_bn3", m, i, o))
resnet50.layer3._modules['0'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_conv1", m, i, o))
resnet50.layer3._modules['0'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_conv2", m, i, o))
resnet50.layer3._modules['0'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block0_conv3", m, i, o))

resnet50.layer3._modules['0'].downsample._modules['0'].register_forward_hook(lambda m, i, o: hook_print("layer3_block0_downsample_conv", m, i, o))
resnet50.layer3._modules['0'].downsample._modules['1'].register_forward_hook(lambda m, i, o: hook_print("layer3_block0_downsample_bn", m, i, o))

resnet50.layer3._modules['1'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_bn1", m, i, o))
resnet50.layer3._modules['1'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_bn2", m, i, o))
resnet50.layer3._modules['1'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_bn3", m, i, o))
resnet50.layer3._modules['1'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_conv1", m, i, o))
resnet50.layer3._modules['1'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_conv2", m, i, o))
resnet50.layer3._modules['1'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block1_conv3", m, i, o))

resnet50.layer3._modules['2'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_bn1", m, i, o))
resnet50.layer3._modules['2'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_bn2", m, i, o))
resnet50.layer3._modules['2'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_bn3", m, i, o))
resnet50.layer3._modules['2'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_conv1", m, i, o))
resnet50.layer3._modules['2'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_conv2", m, i, o))
resnet50.layer3._modules['2'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block2_conv3", m, i, o))

resnet50.layer3._modules['3'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_bn1", m, i, o))
resnet50.layer3._modules['3'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_bn2", m, i, o))
resnet50.layer3._modules['3'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_bn3", m, i, o))
resnet50.layer3._modules['3'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_conv1", m, i, o))
resnet50.layer3._modules['3'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_conv2", m, i, o))
resnet50.layer3._modules['3'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block3_conv3", m, i, o))

resnet50.layer3._modules['4'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_bn1", m, i, o))
resnet50.layer3._modules['4'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_bn2", m, i, o))
resnet50.layer3._modules['4'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_bn3", m, i, o))
resnet50.layer3._modules['4'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_conv1", m, i, o))
resnet50.layer3._modules['4'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_conv2", m, i, o))
resnet50.layer3._modules['4'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block4_conv3", m, i, o))

resnet50.layer3._modules['5'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_bn1", m, i, o))
resnet50.layer3._modules['5'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_bn2", m, i, o))
resnet50.layer3._modules['5'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_bn3", m, i, o))
resnet50.layer3._modules['5'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_conv1", m, i, o))
resnet50.layer3._modules['5'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_conv2", m, i, o))
resnet50.layer3._modules['5'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer3_block5_conv3", m, i, o))



resnet50.layer4._modules['0'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_bn1", m, i, o))
resnet50.layer4._modules['0'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_bn2", m, i, o))
resnet50.layer4._modules['0'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_bn3", m, i, o))
resnet50.layer4._modules['0'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_conv1", m, i, o))
resnet50.layer4._modules['0'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_conv2", m, i, o))
resnet50.layer4._modules['0'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer4_block0_conv3", m, i, o))

resnet50.layer4._modules['0'].downsample._modules['0'].register_forward_hook(lambda m, i, o: hook_print("layer4_block0_downsample_conv", m, i, o))
resnet50.layer4._modules['0'].downsample._modules['1'].register_forward_hook(lambda m, i, o: hook_print("layer4_block0_downsample_bn", m, i, o))

resnet50.layer4._modules['1'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_bn1", m, i, o))
resnet50.layer4._modules['1'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_bn2", m, i, o))
resnet50.layer4._modules['1'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_bn3", m, i, o))
resnet50.layer4._modules['1'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_conv1", m, i, o))
resnet50.layer4._modules['1'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_conv2", m, i, o))
resnet50.layer4._modules['1'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer4_block1_conv3", m, i, o))

resnet50.layer4._modules['2'].bn1.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_bn1", m, i, o))
resnet50.layer4._modules['2'].bn2.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_bn2", m, i, o))
resnet50.layer4._modules['2'].bn3.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_bn3", m, i, o))
resnet50.layer4._modules['2'].conv1.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_conv1", m, i, o))
resnet50.layer4._modules['2'].conv2.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_conv2", m, i, o))
resnet50.layer4._modules['2'].conv3.register_forward_hook(lambda m, i, o: hook_print("layer4_block2_conv3", m, i, o))




# for batch_idx, (data, target) in enumerate(val_loader):
#     out = resnet50(data)
#     break



weightfile.write(strg)
shapefile.write(strgShape)

binaryfile.close()
weightfile.close()


# print(strg)

print("===========================")
print("===========================")
print("===========================")