From e72aa348a0b571e4044a5c67141cce4fc18d462f Mon Sep 17 00:00:00 2001 From: Davide Sapienza Date: Mon, 10 Feb 2020 10:56:33 +0100 Subject: [PATCH] Add DLA34 weights exporter Signed-off-by: Davide Sapienza --- tests/dla34/dla34_weightsexporter.py | 162 +++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 tests/dla34/dla34_weightsexporter.py diff --git a/tests/dla34/dla34_weightsexporter.py b/tests/dla34/dla34_weightsexporter.py new file mode 100644 index 0000000..4a412df --- /dev/null +++ b/tests/dla34/dla34_weightsexporter.py @@ -0,0 +1,162 @@ +import torch +import urllib +from PIL import Image +from torchvision import transforms +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() + fmt = 'f'*len(data) + bin = struct.pack(fmt, *data) + f.write(bin) + +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 exp_input(model, input_batch): + # Export the input batch + model(input_batch) + i = input_batch.cpu().data.numpy() + i = np.array(i, dtype=np.float32) + 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(): + m.eval() + if 'DLAResBlock' in str(m.type): + continue + + in_output = m._value_hook + o = in_output.data.numpy() + 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 + + if ' of Conv2d' in str(m.type) or ' of Linear' in str(m.type): + file_name = "layers/" + t + ".bin" + print("open file: ", file_name) + f = open(file_name, mode='wb') + + 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) + print (" weights shape:", np.shape(w)) + + if 'bias' in m._parameters and m._parameters['bias'] is not None: + b = m._parameters['bias'].data.numpy() + b = np.array(b, dtype=np.float32) + print (" bias shape:", np.shape(b)) + + if 'BatchNorm2d' in str(m.type): + b = m._parameters['bias'].data.numpy() + b = np.array(b, dtype=np.float32) + s = m._parameters['weight'].data.numpy() + s = np.array(s, dtype=np.float32) + rm = m.running_mean.data.numpy() + rm = np.array(rm, dtype=np.float32) + rv = m.running_var.data.numpy() + rv = np.array(rv, dtype=np.float32) + 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: + bin_write(f,w) + if b.size > 0 and b is not None: + bin_write(f,b) + + if ' of BatchNorm2d' in str(m.type) or ' of Linear' in str(m.type): + f.close() + print("close file") + f = None + +if __name__ == '__main__': + model = ptcv_get_model("dla34", pretrained=True) + model.eval() + + # 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) + + # 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) + + # export input bin + exp_input(model, input_batch) + + print_wb_output(model) + + with open("dla34.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))