diff options
Diffstat (limited to 'dnn/torch/weight-exchange')
-rw-r--r-- | dnn/torch/weight-exchange/wexchange/c_export/c_writer.py | 14 | ||||
-rw-r--r-- | dnn/torch/weight-exchange/wexchange/torch/__init__.py | 1 | ||||
-rw-r--r-- | dnn/torch/weight-exchange/wexchange/torch/torch.py | 157 |
3 files changed, 166 insertions, 6 deletions
diff --git a/dnn/torch/weight-exchange/wexchange/c_export/c_writer.py b/dnn/torch/weight-exchange/wexchange/c_export/c_writer.py index 36050881..2745f337 100644 --- a/dnn/torch/weight-exchange/wexchange/c_export/c_writer.py +++ b/dnn/torch/weight-exchange/wexchange/c_export/c_writer.py @@ -38,7 +38,8 @@ class CWriter: create_state_struct=False, enable_binary_blob=True, model_struct_name="Model", - nnet_header="nnet.h"): + nnet_header="nnet.h", + add_typedef=False): """ Writer class for creating souce and header files for weight exports to C @@ -73,6 +74,7 @@ class CWriter: self.enable_binary_blob = enable_binary_blob self.create_state_struct = create_state_struct self.model_struct_name = model_struct_name + self.add_typedef = add_typedef # for binary blob format, format is key=<layer name>, value=(<layer type>, <init call>) self.layer_dict = OrderedDict() @@ -119,11 +121,17 @@ f""" # create model type if self.enable_binary_blob: - self.header.write(f"\nstruct {self.model_struct_name} {{") + if self.add_typedef: + self.header.write(f"\ntypedef struct {{") + else: + self.header.write(f"\nstruct {self.model_struct_name} {{") for name, data in self.layer_dict.items(): layer_type = data[0] self.header.write(f"\n {layer_type} {name};") - self.header.write(f"\n}};\n") + if self.add_typedef: + self.header.write(f"\n}} {self.model_struct_name};\n") + else: + self.header.write(f"\n}};\n") init_prototype = f"int init_{self.model_struct_name.lower()}({self.model_struct_name} *model, const WeightArray *arrays)" self.header.write(f"\n{init_prototype};\n") diff --git a/dnn/torch/weight-exchange/wexchange/torch/__init__.py b/dnn/torch/weight-exchange/wexchange/torch/__init__.py index 98c96fad..8245566d 100644 --- a/dnn/torch/weight-exchange/wexchange/torch/__init__.py +++ b/dnn/torch/weight-exchange/wexchange/torch/__init__.py @@ -34,3 +34,4 @@ from .torch import dump_torch_gru_weights, load_torch_gru_weights from .torch import dump_torch_grucell_weights from .torch import dump_torch_embedding_weights, load_torch_embedding_weights from .torch import dump_torch_weights, load_torch_weights +from .torch import dump_torch_adaptive_conv1d_weights
\ No newline at end of file diff --git a/dnn/torch/weight-exchange/wexchange/torch/torch.py b/dnn/torch/weight-exchange/wexchange/torch/torch.py index 281d9be3..f7e16032 100644 --- a/dnn/torch/weight-exchange/wexchange/torch/torch.py +++ b/dnn/torch/weight-exchange/wexchange/torch/torch.py @@ -28,12 +28,154 @@ """ import os +import sys import torch import numpy as np +sys.path.append(sys.path.append(os.path.join(os.path.dirname(__file__), '../osce'))) +try: + import utils.layers as osce_layers + from utils.layers.limited_adaptive_conv1d import LimitedAdaptiveConv1d + from utils.layers.limited_adaptive_comb1d import LimitedAdaptiveComb1d + from utils.layers.td_shaper import TDShaper + has_osce=True +except: + has_osce=False + from wexchange.c_export import CWriter, print_gru_layer, print_dense_layer, print_conv1d_layer, print_tconv1d_layer, print_conv2d_layer +def dump_torch_adaptive_conv1d_weights(where, adaconv, name='adaconv', scale=1/128, quantize=False): + + + w_kernel = adaconv.conv_kernel.weight.detach().cpu().numpy().copy() + b_kernel = adaconv.conv_kernel.bias.detach().cpu().numpy().copy() + w_gain = adaconv.filter_gain.weight.detach().cpu().numpy().copy() + b_gain = adaconv.filter_gain.bias.detach().cpu().numpy().copy() + + if isinstance(where, CWriter): + # pad kernel for quantization + left_padding = adaconv.padding[0] + kernel_size = adaconv.kernel_size + in_channels = adaconv.in_channels + out_channels = adaconv.out_channels + feature_dim = adaconv.feature_dim + + if quantize and kernel_size % 8: + kernel_padding = 8 - (kernel_size % 8) + w_kernel = np.concatenate( + (np.zeros((out_channels, in_channels, kernel_padding, feature_dim)), w_kernel.reshape(out_channels, in_channels, kernel_size, feature_dim)), + dtype=w_kernel.dtype, + axis=2).reshape(-1, feature_dim) + b_kernel = np.concatenate( + (np.zeros((out_channels, in_channels, kernel_padding)), b_kernel.reshape(out_channels, in_channels, kernel_size)), + dtype=b_kernel.dtype, + axis=2).reshape(-1) + left_padding += kernel_padding + kernel_size += kernel_padding + + # write relevant scalar parameters to header file + where.header.write(f""" +#define {name.upper()}_FILTER_GAIN_A {adaconv.filter_gain_a:f}f +#define {name.upper()}_FILTER_GAIN_B {adaconv.filter_gain_b:f}f +#define {name.upper()}_SHAPE_GAIN {adaconv.shape_gain:f}f +#define {name.upper()}_KERNEL_SIZE {kernel_size} +#define {name.upper()}_FRAME_SIZE {adaconv.frame_size} +#define {name.upper()}_LEFT_PADDING {left_padding} +#define {name.upper()}_OVERLAP_SIZE {adaconv.overlap_size} +#define {name.upper()}_IN_CHANNELS {adaconv.in_channels} +#define {name.upper()}_OUT_CHANNELS {adaconv.out_channels} +#define {name.upper()}_NORM_P {adaconv.norm_p} +#define {name.upper()}_FEATURE_DIM {adaconv.feature_dim} +""" + ) + + print_dense_layer(where, name + "_kernel", w_kernel, b_kernel, scale=scale, format='torch', sparse=False, diagonal=False, quantize=quantize) + print_dense_layer(where, name + "_gain", w_gain, b_gain, format='torch', sparse=False, diagonal=False, quantize=False) + + + else: + np.save(where, 'weight_kernel.npy', w_kernel) + np.save(where, 'bias_kernel.npy', b_kernel) + np.save(where, 'weight_gain.npy', w_gain) + np.save(where, 'bias_gain.npy', b_gain) + + +def dump_torch_adaptive_comb1d_weights(where, adaconv, name='adaconv', scale=1/128, quantize=False): + + + w_kernel = adaconv.conv_kernel.weight.detach().cpu().numpy().copy() + b_kernel = adaconv.conv_kernel.bias.detach().cpu().numpy().copy() + w_gain = adaconv.filter_gain.weight.detach().cpu().numpy().copy() + b_gain = adaconv.filter_gain.bias.detach().cpu().numpy().copy() + w_global_gain = adaconv.global_filter_gain.weight.detach().cpu().numpy().copy() + b_global_gain = adaconv.global_filter_gain.bias.detach().cpu().numpy().copy() + + + if isinstance(where, CWriter): + # pad kernel for quantization + left_padding = adaconv.padding[0] + kernel_size = adaconv.kernel_size + + if quantize and w_kernel.shape[0] % 8: + kernel_padding = 8 - (w_kernel.shape[0] % 8) + w_kernel = np.concatenate((np.zeros((kernel_padding, w_kernel.shape[1])), w_kernel), dtype=w_kernel.dtype) + b_kernel = np.concatenate((np.zeros((kernel_padding)), b_kernel), dtype=b_kernel.dtype) + left_padding += kernel_padding + kernel_size += kernel_padding + # write relevant scalar parameters to header file + where.header.write(f""" +#define {name.upper()}_FILTER_GAIN_A {adaconv.filter_gain_a:f}f +#define {name.upper()}_FILTER_GAIN_B {adaconv.filter_gain_b:f}f +#define {name.upper()}_LOG_GAIN_LIMIT {adaconv.log_gain_limit:f}f +#define {name.upper()}_KERNEL_SIZE {kernel_size} +#define {name.upper()}_LEFT_PADDING {left_padding} +#define {name.upper()}_FRAME_SIZE {adaconv.frame_size} +#define {name.upper()}_OVERLAP_SIZE {adaconv.overlap_size} +#define {name.upper()}_IN_CHANNELS {adaconv.in_channels} +#define {name.upper()}_OUT_CHANNELS {adaconv.out_channels} +#define {name.upper()}_NORM_P {adaconv.norm_p} +#define {name.upper()}_FEATURE_DIM {adaconv.feature_dim} +#define {name.upper()}_MAX_LAG {adaconv.max_lag} +""" + ) + + print_dense_layer(where, name + "_kernel", w_kernel, b_kernel, scale=scale, format='torch', sparse=False, diagonal=False, quantize=quantize) + print_dense_layer(where, name + "_gain", w_gain, b_gain, format='torch', sparse=False, diagonal=False, quantize=False) + print_dense_layer(where, name + "_global_gain", w_global_gain, b_global_gain, format='torch', sparse=False, diagonal=False, quantize=False) + + + else: + np.save(where, 'weight_kernel.npy', w_kernel) + np.save(where, 'bias_kernel.npy', b_kernel) + np.save(where, 'weight_gain.npy', w_gain) + np.save(where, 'bias_gain.npy', b_gain) + np.save(where, 'weight_global_gain.npy', w_global_gain) + np.save(where, 'bias_global_gain.npy', b_global_gain) + +def dump_torch_tdshaper(where, shaper, name='tdshaper'): + + if isinstance(where, CWriter): + where.header.write(f""" +#define {name.upper()}_FEATURE_DIM {shaper.feature_dim} +#define {name.upper()}_FRAME_SIZE {shaper.frame_size} +#define {name.upper()}_AVG_POOL_K {shaper.avg_pool_k} +#define {name.upper()}_INNOVATE {1 if shaper.innovate else 0} +#define {name.upper()}_POOL_AFTER {1 if shaper.pool_after else 0} +""" + ) + + dump_torch_conv1d_weights(where, shaper.feature_alpha1, name + "_alpha1") + dump_torch_conv1d_weights(where, shaper.feature_alpha2, name + "_alpha2") + + if shaper.innovate: + dump_torch_conv1d_weights(where, shaper.feature_alpha1b, name + "_alpha1b") + dump_torch_conv1d_weights(where, shaper.feature_alpha1c, name + "_alpha1c") + dump_torch_conv1d_weights(where, shaper.feature_alpha2b, name + "_alpha2b") + dump_torch_conv1d_weights(where, shaper.feature_alpha2c, name + "_alpha2c") + + + def dump_torch_gru_weights(where, gru, name='gru', input_sparse=False, recurrent_sparse=False, quantize=False, scale=1/128, recurrent_scale=1/128): assert gru.num_layers == 1 @@ -221,7 +363,6 @@ def load_torch_conv2d_weights(where, conv): def dump_torch_embedding_weights(where, embed, name='embed', scale=1/128, sparse=False, diagonal=False, quantize=False): - print("quantize = ", quantize) w = embed.weight.detach().cpu().numpy().copy().transpose() b = np.zeros(w.shape[0], dtype=w.dtype) @@ -257,11 +398,21 @@ def dump_torch_weights(where, module, name=None, verbose=False, **kwargs): elif isinstance(module, torch.nn.Conv2d): return dump_torch_conv2d_weights(where, module, name, **kwargs) elif isinstance(module, torch.nn.Embedding): - return dump_torch_embedding_weights(where, module) + return dump_torch_embedding_weights(where, module, name, **kwargs) elif isinstance(module, torch.nn.ConvTranspose1d): return dump_torch_tconv1d_weights(where, module, name, **kwargs) else: - raise ValueError(f'dump_torch_weights: layer of type {type(module)} not supported') + if has_osce: + if isinstance(module, LimitedAdaptiveConv1d): + dump_torch_adaptive_conv1d_weights(where, module, name, **kwargs) + elif isinstance(module, LimitedAdaptiveComb1d): + dump_torch_adaptive_comb1d_weights(where, module, name, **kwargs) + elif isinstance(module, TDShaper): + dump_torch_tdshaper(where, module, name, **kwargs) + else: + raise ValueError(f'dump_torch_weights: layer of type {type(module)} not supported') + else: + raise ValueError(f'dump_torch_weights: layer of type {type(module)} not supported') def load_torch_weights(where, module): """ generic function for loading weights of some torch.nn.Module """ |