Add DLA34 weights exporter
Signed-off-by: Davide Sapienza <sapienza.dav@gmail.com>
This commit is contained in:
Davide Sapienza
@@ -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))
|
||||
Reference in New Issue
Block a user