mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Update ResNet101 weights exporter

Browse Source 
Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
This commit is contained in:
Davide Sapienza
2020-02-10 10:54:04 +01:00
parent 9007e25a00
commit 62fe82ce9e
1 changed files with 69 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/resnet101/resnet101_weightsexporter.py
+69 -70
View File
@@ -2,15 +2,26 @@ import torch
import urllib import urllib
from PIL import Image from PIL import Image
from torchvision import transforms from torchvision import transforms
from torchsummary import summary
import numpy as np import numpy as np
import struct 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 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): def bin_write(f, data):
data =data.flatten() data =data.flatten()
# print(data)
fmt = 'f'*len(data) fmt = 'f'*len(data)
bin = struct.pack(fmt, *data) bin = struct.pack(fmt, *data)
f.write(bin) f.write(bin)
@@ -18,38 +29,46 @@ def bin_write(f, data):
def hook(module, input, output): def hook(module, input, output):
setattr(module, "_value_hook", 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): def exp_input(model, input_batch):
for name, param in model.named_parameters(): # Export the input batch
# 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)
model(input_batch) model(input_batch)
i = input_batch.data.numpy() i = input_batch.cpu().data.numpy()
i = np.array(i, dtype=np.float32) i = np.array(i, dtype=np.float32)
print(i.shape)
i.tofile("debug/input.bin", format="f") i.tofile("debug/input.bin", format="f")
print("input: ", i.shape)
def print_wb_output(model):
f = None f = None
for n, m in model.named_modules(): for n, m in model.named_modules():
in_output = m._value_hook in_output = m._value_hook
@@ -57,7 +76,9 @@ def print_wb_output(model, input_batch):
o = np.array(o, dtype=np.float32) o = np.array(o, dtype=np.float32)
t = '-'.join(n.split('.')) t = '-'.join(n.split('.'))
o.tofile("debug/" + t + ".bin", format="f") 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)): if not(' of Conv2d' in str(m.type) or ' of Linear' in str(m.type) or ' of BatchNorm2d' in str(m.type)):
continue continue
@@ -66,14 +87,8 @@ def print_wb_output(model, input_batch):
print("open file: ", file_name) print("open file: ", file_name)
f = open(file_name, mode='wb') f = open(file_name, mode='wb')
print(n, ' ----------------------------------------------------------------')
# print(m._parameters)
#print(m.type)
w = np.array([]) w = np.array([])
b = np.array([]) b = np.array([])
if 'weight' in m._parameters and m._parameters['weight'] is not None: if 'weight' in m._parameters and m._parameters['weight'] is not None:
w = m._parameters['weight'].data.numpy() w = m._parameters['weight'].data.numpy()
w = np.array(w, dtype=np.float32) 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 = m._parameters['bias'].data.numpy()
b = np.array(b, dtype=np.float32) b = np.array(b, dtype=np.float32)
print (" bias shape:", np.shape(b)) 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): if 'BatchNorm2d' in str(m.type):
b = m._parameters['bias'].data.numpy() b = m._parameters['bias'].data.numpy()
@@ -96,30 +108,24 @@ def print_wb_output(model, input_batch):
rm = np.array(rm, dtype=np.float32) rm = np.array(rm, dtype=np.float32)
rv = m.running_var.data.numpy() rv = m.running_var.data.numpy()
rv = np.array(rv, dtype=np.float32) rv = np.array(rv, dtype=np.float32)
#s.tofile(f, format="f")
bin_write(f,b) bin_write(f,b)
bin_write(f,s) bin_write(f,s)
bin_write(f,rm) bin_write(f,rm)
bin_write(f,rv) bin_write(f,rv)
print (" b shape:", np.shape(b))
print (" s shape:", np.shape(s)) print (" s shape:", np.shape(s))
print (" rm shape:", np.shape(rm)) print (" rm shape:", np.shape(rm))
print (" rv shape:", np.shape(rv)) print (" rv shape:", np.shape(rv))
else: else:
# w.tofile(f, format="f")
bin_write(f,w) bin_write(f,w)
# print("w- ",w)
if b.size > 0: if b.size > 0:
# b.tofile(f, format="f")
bin_write(f,b) bin_write(f,b)
# print("b - ",b)
if ' of BatchNorm2d' in str(m.type) or ' of Linear' in str(m.type): if ' of BatchNorm2d' in str(m.type) or ' of Linear' in str(m.type):
f.close() f.close()
print("close file") print("close file")
f = None f = None
# return
@@ -130,34 +136,27 @@ if __name__ == '__main__':
model = torch.hub.load('pytorch/vision', 'resnet101', pretrained=True) model = torch.hub.load('pytorch/vision', 'resnet101', pretrained=True)
model.eval() model.eval()
# Download an example image from the pytorch website # load an example image and load it on model
url, filename = ("https://github.com/pytorch/hub/raw/master/dog.jpg", "dog.jpg") model, input_batch = load_ex_image(model)
try: urllib.URLopener().retrieve(url, filename) model.eval()
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')
with torch.no_grad(): with torch.no_grad():
output = model(input_batch) output = model(input_batch)
# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes # create folders debug and layers if do not exist
# print(output) 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))