From 62fe82ce9e905cf815eb74b02b4b5ffd47d0253e Mon Sep 17 00:00:00 2001 From: Davide Sapienza Date: Mon, 10 Feb 2020 10:54:04 +0100 Subject: [PATCH] Update ResNet101 weights exporter Signed-off-by: Davide Sapienza --- tests/resnet101/resnet101_weightsexporter.py | 139 +++++++++---------- 1 file changed, 69 insertions(+), 70 deletions(-) diff --git a/tests/resnet101/resnet101_weightsexporter.py b/tests/resnet101/resnet101_weightsexporter.py index 140154e..e10a037 100644 --- a/tests/resnet101/resnet101_weightsexporter.py +++ b/tests/resnet101/resnet101_weightsexporter.py @@ -2,15 +2,26 @@ import torch import urllib from PIL import Image from torchvision import transforms -from torchsummary import summary import numpy as np import struct +import os +from pytorchcv.model_provider import get_model as ptcv_get_model +from torch.autograd import Variable + +from torchsummary import summary import torch.nn as nn +from torch.jit import trace + +def create_folders(): + if not os.path.exists('debug'): + os.makedirs('debug') + if not os.path.exists('layers'): + os.makedirs('layers') + def bin_write(f, data): data =data.flatten() - # print(data) fmt = 'f'*len(data) bin = struct.pack(fmt, *data) f.write(bin) @@ -18,38 +29,46 @@ def bin_write(f, data): def hook(module, input, output): setattr(module, "_value_hook", output) +def load_ex_image(model): + # Download an example image from the pytorch website + url, filename = ( + "https://github.com/pytorch/hub/raw/master/dog.jpg", "dog.jpg") + try: + urllib.URLopener().retrieve(url, filename) + except: + urllib.request.urlretrieve(url, filename) + + # sample execution (requires torchvision) + input_image = Image.open(filename) + print("input_image: ",input_image.size) + preprocess = transforms.Compose([ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[ + 0.229, 0.224, 0.225]), + ]) + input_tensor = preprocess(input_image) + print("input_tensor: ",input_tensor.shape) + # create a mini-batch as expected by the model + input_batch = input_tensor.unsqueeze(0) + # move the input and model to GPU for speed if available + if torch.cuda.is_available(): + input_batch = input_batch.to('cuda') + model.to('cuda') + + return model, input_batch -def print_wb(model, folder): - for name, param in model.named_parameters(): - # print ("Layer", name) - t = name.split('.')[0:-1] - arg = name.split('.')[-1] - t = '-'.join(t) - print (" type: ", t) - - if arg == 'weight': - w = param.data.numpy() - print (" weights shape:", np.shape(w)) - w.tofile(folder + "/" + t + ".bin", format="f") - elif arg == 'bias': - b = param.data.numpy() - print (" bias shape:", np.shape(b)) - b.tofile(folder + "/" + t + ".bias.bin", format="f") - else: - print("Ops!") - - -def print_wb_output(model, input_batch): - for n, m in model.named_modules(): - m.register_forward_hook(hook) - +def exp_input(model, input_batch): + # Export the input batch model(input_batch) - i = input_batch.data.numpy() + i = input_batch.cpu().data.numpy() i = np.array(i, dtype=np.float32) - print(i.shape) i.tofile("debug/input.bin", format="f") + print("input: ", i.shape) +def print_wb_output(model): f = None for n, m in model.named_modules(): in_output = m._value_hook @@ -57,7 +76,9 @@ def print_wb_output(model, input_batch): o = np.array(o, dtype=np.float32) t = '-'.join(n.split('.')) o.tofile("debug/" + t + ".bin", format="f") - + print('------- ', n, ' ------') + print("debug ",o.shape) + if not(' of Conv2d' in str(m.type) or ' of Linear' in str(m.type) or ' of BatchNorm2d' in str(m.type)): continue @@ -66,14 +87,8 @@ def print_wb_output(model, input_batch): print("open file: ", file_name) f = open(file_name, mode='wb') - print(n, ' ----------------------------------------------------------------') - # print(m._parameters) - #print(m.type) - w = np.array([]) b = np.array([]) - - if 'weight' in m._parameters and m._parameters['weight'] is not None: w = m._parameters['weight'].data.numpy() w = np.array(w, dtype=np.float32) @@ -83,9 +98,6 @@ def print_wb_output(model, input_batch): b = m._parameters['bias'].data.numpy() b = np.array(b, dtype=np.float32) print (" bias shape:", np.shape(b)) - # else: - # b = np.zeros(w.shape[0], dtype=np.float32) - # print (" bias shape:", np.shape(b)) if 'BatchNorm2d' in str(m.type): b = m._parameters['bias'].data.numpy() @@ -96,30 +108,24 @@ def print_wb_output(model, input_batch): rm = np.array(rm, dtype=np.float32) rv = m.running_var.data.numpy() rv = np.array(rv, dtype=np.float32) - #s.tofile(f, format="f") bin_write(f,b) bin_write(f,s) bin_write(f,rm) bin_write(f,rv) + print (" b shape:", np.shape(b)) print (" s shape:", np.shape(s)) print (" rm shape:", np.shape(rm)) print (" rv shape:", np.shape(rv)) else: - - # w.tofile(f, format="f") bin_write(f,w) - # print("w- ",w) if b.size > 0: - # b.tofile(f, format="f") bin_write(f,b) - # print("b - ",b) if ' of BatchNorm2d' in str(m.type) or ' of Linear' in str(m.type): f.close() print("close file") f = None - # return @@ -130,34 +136,27 @@ if __name__ == '__main__': model = torch.hub.load('pytorch/vision', 'resnet101', pretrained=True) model.eval() - # Download an example image from the pytorch website - url, filename = ("https://github.com/pytorch/hub/raw/master/dog.jpg", "dog.jpg") - try: urllib.URLopener().retrieve(url, filename) - except: urllib.request.urlretrieve(url, filename) - - # sample execution (requires torchvision) - input_image = Image.open(filename) - preprocess = transforms.Compose([ - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ]) - input_tensor = preprocess(input_image) - input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model - - # move the input and model to GPU for speed if available - if torch.cuda.is_available(): - input_batch = input_batch.to('cuda') - model.to('cuda') - + # load an example image and load it on model + model, input_batch = load_ex_image(model) + model.eval() with torch.no_grad(): output = model(input_batch) - # Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes - # print(output) + # create folders debug and layers if do not exist + create_folders() + # add output attribute to the layers + for n, m in model.named_modules(): + m.register_forward_hook(hook) - print_wb_output(model, input_batch) + # export input bin + exp_input(model, input_batch) - # print(list(model.children())) + print_wb_output(model) + + with open("resnet101.txt", 'w') as f: + for item in list(model.children()): + f.write("%s\n" % item) + + summary(model, (3, 224, 224)) + # print(trace(model, input_batch))