diff options
| author | Jan Buethe <jbuethe@amazon.de> | 2023-10-20 15:14:31 +0300 |
|---|---|---|
| committer | Jan Buethe <jbuethe@amazon.de> | 2023-10-20 15:14:31 +0300 |
| commit | 1accd2472e678d540fa024f05da68088014dafaa (patch) | |
| tree | 7032c019b4b6e2ef48dbe739dad0284845ff0cf6 | |
| parent | 88c8b3078518b649933616fb7c9a78e4d086233a (diff) | |
finalized quantization option in export_rdovae_weights.py
| -rw-r--r-- | dnn/torch/rdovae/export_rdovae_weights.py | 20 |
1 files changed, 10 insertions, 10 deletions
diff --git a/dnn/torch/rdovae/export_rdovae_weights.py b/dnn/torch/rdovae/export_rdovae_weights.py index 856ec93e..9a35c17a 100644 --- a/dnn/torch/rdovae/export_rdovae_weights.py +++ b/dnn/torch/rdovae/export_rdovae_weights.py @@ -121,9 +121,9 @@ f""" ('core_encoder.module.state_dense_2' , 'gdense2' , 'TANH', True) ] - for name, export_name, _, _ in encoder_dense_layers: + for name, export_name, _, quantize in encoder_dense_layers: layer = model.get_submodule(name) - dump_torch_weights(enc_writer, layer, name=export_name, verbose=True) + dump_torch_weights(enc_writer, layer, name=export_name, verbose=True, quantize=quantize, scale=None) encoder_gru_layers = [ @@ -134,8 +134,8 @@ f""" ('core_encoder.module.gru5' , 'enc_gru5', 'TANH', True), ] - enc_max_rnn_units = max([dump_torch_weights(enc_writer, model.get_submodule(name), export_name, verbose=True, input_sparse=True, quantize=True) - for name, export_name, _, _ in encoder_gru_layers]) + enc_max_rnn_units = max([dump_torch_weights(enc_writer, model.get_submodule(name), export_name, verbose=True, input_sparse=True, quantize=quantize, scale=None, recurrent_scale=None) + for name, export_name, _, quantize in encoder_gru_layers]) encoder_conv_layers = [ @@ -146,7 +146,7 @@ f""" ('core_encoder.module.conv5.conv' , 'enc_conv5', 'TANH', True), ] - enc_max_conv_inputs = max([dump_torch_weights(enc_writer, model.get_submodule(name), export_name, verbose=True, quantize=False) for name, export_name, _, _ in encoder_conv_layers]) + enc_max_conv_inputs = max([dump_torch_weights(enc_writer, model.get_submodule(name), export_name, verbose=True, quantize=quantize, scale=None) for name, export_name, _, quantize in encoder_conv_layers]) del enc_writer @@ -159,9 +159,9 @@ f""" ('core_decoder.module.gru_init' , 'dec_gru_init', 'TANH', True), ] - for name, export_name, _, _ in decoder_dense_layers: + for name, export_name, _, quantize in decoder_dense_layers: layer = model.get_submodule(name) - dump_torch_weights(dec_writer, layer, name=export_name, verbose=True) + dump_torch_weights(dec_writer, layer, name=export_name, verbose=True, quantize=quantize, scale=None) decoder_gru_layers = [ @@ -172,8 +172,8 @@ f""" ('core_decoder.module.gru5' , 'dec_gru5', 'TANH', True), ] - dec_max_rnn_units = max([dump_torch_weights(dec_writer, model.get_submodule(name), export_name, verbose=True, input_sparse=True, quantize=True) - for name, export_name, _, _ in decoder_gru_layers]) + dec_max_rnn_units = max([dump_torch_weights(dec_writer, model.get_submodule(name), export_name, verbose=True, input_sparse=True, quantize=quantize, scale=None, recurrent_scale=None) + for name, export_name, _, quantize in decoder_gru_layers]) decoder_conv_layers = [ ('core_decoder.module.conv1.conv' , 'dec_conv1', 'TANH', True), @@ -183,7 +183,7 @@ f""" ('core_decoder.module.conv5.conv' , 'dec_conv5', 'TANH', True), ] - dec_max_conv_inputs = max([dump_torch_weights(dec_writer, model.get_submodule(name), export_name, verbose=True, quantize=False) for name, export_name, _, _ in decoder_conv_layers]) + dec_max_conv_inputs = max([dump_torch_weights(dec_writer, model.get_submodule(name), export_name, verbose=True, quantize=quantize, scale=None) for name, export_name, _, quantize in decoder_conv_layers]) del dec_writer |