Skip to content

Commit

Permalink
Marking RDOVAE layers to quantize
Browse files Browse the repository at this point in the history
  • Loading branch information
jmvalin committed Oct 19, 2023
1 parent 60ac1c6 commit d720955
Showing 1 changed file with 25 additions and 25 deletions.
50 changes: 25 additions & 25 deletions dnn/torch/rdovae/export_rdovae_weights.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,9 @@ def c_export(args, model):
# encoder
encoder_dense_layers = [
('core_encoder.module.dense_1' , 'enc_dense1', 'TANH', False,),
('core_encoder.module.z_dense' , 'enc_zdense', 'LINEAR', False,),
('core_encoder.module.state_dense_1' , 'gdense1' , 'TANH', False,),
('core_encoder.module.state_dense_2' , 'gdense2' , 'TANH', False)
('core_encoder.module.z_dense' , 'enc_zdense', 'LINEAR', True,),
('core_encoder.module.state_dense_1' , 'gdense1' , 'TANH', True,),
('core_encoder.module.state_dense_2' , 'gdense2' , 'TANH', True)
]

for name, export_name, _, _ in encoder_dense_layers:
Expand All @@ -127,23 +127,23 @@ def c_export(args, model):


encoder_gru_layers = [
('core_encoder.module.gru1' , 'enc_gru1', 'TANH', False),
('core_encoder.module.gru2' , 'enc_gru2', 'TANH', False),
('core_encoder.module.gru3' , 'enc_gru3', 'TANH', False),
('core_encoder.module.gru4' , 'enc_gru4', 'TANH', False),
('core_encoder.module.gru5' , 'enc_gru5', 'TANH', False),
('core_encoder.module.gru1' , 'enc_gru1', 'TANH', True),
('core_encoder.module.gru2' , 'enc_gru2', 'TANH', True),
('core_encoder.module.gru3' , 'enc_gru3', 'TANH', True),
('core_encoder.module.gru4' , 'enc_gru4', 'TANH', True),
('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])


encoder_conv_layers = [
('core_encoder.module.conv1.conv' , 'enc_conv1', 'TANH', False),
('core_encoder.module.conv2.conv' , 'enc_conv2', 'TANH', False),
('core_encoder.module.conv3.conv' , 'enc_conv3', 'TANH', False),
('core_encoder.module.conv4.conv' , 'enc_conv4', 'TANH', False),
('core_encoder.module.conv5.conv' , 'enc_conv5', 'TANH', False),
('core_encoder.module.conv1.conv' , 'enc_conv1', 'TANH', True),
('core_encoder.module.conv2.conv' , 'enc_conv2', 'TANH', True),
('core_encoder.module.conv3.conv' , 'enc_conv3', 'TANH', True),
('core_encoder.module.conv4.conv' , 'enc_conv4', 'TANH', True),
('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])
Expand All @@ -154,9 +154,9 @@ def c_export(args, model):
# decoder
decoder_dense_layers = [
('core_decoder.module.dense_1' , 'dec_dense1', 'TANH', False),
('core_decoder.module.output' , 'dec_output', 'LINEAR', False),
('core_decoder.module.output' , 'dec_output', 'LINEAR', True),
('core_decoder.module.hidden_init' , 'dec_hidden_init', 'TANH', False),
('core_decoder.module.gru_init' , 'dec_gru_init', 'TANH', False),
('core_decoder.module.gru_init' , 'dec_gru_init', 'TANH', True),
]

for name, export_name, _, _ in decoder_dense_layers:
Expand All @@ -165,22 +165,22 @@ def c_export(args, model):


decoder_gru_layers = [
('core_decoder.module.gru1' , 'dec_gru1', 'TANH', False),
('core_decoder.module.gru2' , 'dec_gru2', 'TANH', False),
('core_decoder.module.gru3' , 'dec_gru3', 'TANH', False),
('core_decoder.module.gru4' , 'dec_gru4', 'TANH', False),
('core_decoder.module.gru5' , 'dec_gru5', 'TANH', False),
('core_decoder.module.gru1' , 'dec_gru1', 'TANH', True),
('core_decoder.module.gru2' , 'dec_gru2', 'TANH', True),
('core_decoder.module.gru3' , 'dec_gru3', 'TANH', True),
('core_decoder.module.gru4' , 'dec_gru4', 'TANH', True),
('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])

decoder_conv_layers = [
('core_decoder.module.conv1.conv' , 'dec_conv1', 'TANH', False),
('core_decoder.module.conv2.conv' , 'dec_conv2', 'TANH', False),
('core_decoder.module.conv3.conv' , 'dec_conv3', 'TANH', False),
('core_decoder.module.conv4.conv' , 'dec_conv4', 'TANH', False),
('core_decoder.module.conv5.conv' , 'dec_conv5', 'TANH', False),
('core_decoder.module.conv1.conv' , 'dec_conv1', 'TANH', True),
('core_decoder.module.conv2.conv' , 'dec_conv2', 'TANH', True),
('core_decoder.module.conv3.conv' , 'dec_conv3', 'TANH', True),
('core_decoder.module.conv4.conv' , 'dec_conv4', 'TANH', True),
('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])
Expand Down

0 comments on commit d720955

Please sign in to comment.