@scrypted/libav 1.0.167 → 1.0.168

package/opus/dnn/torch/plc/train_plc.py

@@ -0,0 +1,145 @@
+import os
+import argparse
+import random
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import tqdm
+
+import plc
+from plc_dataset import PLCDataset
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('features', type=str, help='path to feature file in .f32 format')
+parser.add_argument('loss', type=str, help='path to signal file in .s8 format')
+parser.add_argument('output', type=str, help='path to output folder')
+
+parser.add_argument('--suffix', type=str, help="model name suffix", default="")
+parser.add_argument('--cuda-visible-devices', type=str, help="comma separates list of cuda visible device indices, default: CUDA_VISIBLE_DEVICES", default=None)
+
+
+model_group = parser.add_argument_group(title="model parameters")
+model_group.add_argument('--cond-size', type=int, help="first conditioning size, default: 128", default=128)
+model_group.add_argument('--gru-size', type=int, help="GRU size, default: 128", default=128)
+
+training_group = parser.add_argument_group(title="training parameters")
+training_group.add_argument('--batch-size', type=int, help="batch size, default: 512", default=512)
+training_group.add_argument('--lr', type=float, help='learning rate, default: 1e-3', default=1e-3)
+training_group.add_argument('--epochs', type=int, help='number of training epochs, default: 20', default=20)
+training_group.add_argument('--sequence-length', type=int, help='sequence length, default: 15', default=15)
+training_group.add_argument('--lr-decay', type=float, help='learning rate decay factor, default: 1e-4', default=1e-4)
+training_group.add_argument('--initial-checkpoint', type=str, help='initial checkpoint to start training from, default: None', default=None)
+
+args = parser.parse_args()
+
+if args.cuda_visible_devices != None:
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda_visible_devices
+
+# checkpoints
+checkpoint_dir = os.path.join(args.output, 'checkpoints')
+checkpoint = dict()
+os.makedirs(checkpoint_dir, exist_ok=True)
+
+
+# training parameters
+batch_size = args.batch_size
+lr = args.lr
+epochs = args.epochs
+sequence_length = args.sequence_length
+lr_decay = args.lr_decay
+
+adam_betas = [0.8, 0.95]
+adam_eps = 1e-8
+features_file = args.features
+loss_file = args.loss
+
+# model parameters
+cond_size  = args.cond_size
+
+
+checkpoint['batch_size'] = batch_size
+checkpoint['lr'] = lr
+checkpoint['lr_decay'] = lr_decay
+checkpoint['epochs'] = epochs
+checkpoint['sequence_length'] = sequence_length
+checkpoint['adam_betas'] = adam_betas
+
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+checkpoint['model_args']    = ()
+checkpoint['model_kwargs']  = {'cond_size': cond_size, 'gru_size': args.gru_size}
+print(checkpoint['model_kwargs'])
+model = plc.PLC(*checkpoint['model_args'], **checkpoint['model_kwargs'])
+
+if type(args.initial_checkpoint) != type(None):
+    checkpoint = torch.load(args.initial_checkpoint, map_location='cpu')
+    model.load_state_dict(checkpoint['state_dict'], strict=False)
+
+checkpoint['state_dict']    = model.state_dict()
+
+
+dataset = PLCDataset(features_file, loss_file, sequence_length=sequence_length)
+dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=4)
+
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=adam_betas, eps=adam_eps)
+
+
+# learning rate scheduler
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lambda x : 1 / (1 + lr_decay * x))
+
+states = None
+
+plc_loss = plc.plc_loss(18, device=device)
+if __name__ == '__main__':
+    model.to(device)
+
+    for epoch in range(1, epochs + 1):
+
+        running_loss = 0
+        running_l1_loss = 0
+        running_ceps_loss = 0
+        running_band_loss = 0
+        running_pitch_loss = 0
+
+        print(f"training epoch {epoch}...")
+        with tqdm.tqdm(dataloader, unit='batch') as tepoch:
+            for i, (features, lost, target) in enumerate(tepoch):
+                optimizer.zero_grad()
+                features = features.to(device)
+                lost = lost.to(device)
+                target = target.to(device)
+
+                out, states = model(features, lost)
+
+                loss, l1_loss, ceps_loss, band_loss, pitch_loss = plc_loss(target, out)
+
+                loss.backward()
+                optimizer.step()
+
+                #model.clip_weights()
+
+                scheduler.step()
+
+                running_loss += loss.detach().cpu().item()
+                running_l1_loss += l1_loss.detach().cpu().item()
+                running_ceps_loss += ceps_loss.detach().cpu().item()
+                running_band_loss += band_loss.detach().cpu().item()
+                running_pitch_loss += pitch_loss.detach().cpu().item()
+                tepoch.set_postfix(loss=f"{running_loss/(i+1):8.5f}",
+                                   l1_loss=f"{running_l1_loss/(i+1):8.5f}",
+                                   ceps_loss=f"{running_ceps_loss/(i+1):8.5f}",
+                                   band_loss=f"{running_band_loss/(i+1):8.5f}",
+                                   pitch_loss=f"{running_pitch_loss/(i+1):8.5f}",
+                                   )
+
+        # save checkpoint
+        checkpoint_path = os.path.join(checkpoint_dir, f'plc{args.suffix}_{epoch}.pth')
+        checkpoint['state_dict'] = model.state_dict()
+        checkpoint['loss'] = running_loss / len(dataloader)
+        checkpoint['epoch'] = epoch
+        torch.save(checkpoint, checkpoint_path)

package/opus/dnn/torch/rdovae/README.md

@@ -0,0 +1,55 @@
+# Deep REDundancy (DRED) with RDO-VAE
+
+This is a rate-distortion-optimized variational autoencoder (RDO-VAE) designed
+to coding redundancy information. Pre-trained models are provided as C code
+in the dnn/ directory with the corresponding model in dnn/models/ directory
+(name starts with rdovae_). If you don't want to train a new DRED model, you can
+skip straight to the Inference section.
+
+## Data preparation
+
+First, fetch all the data from the datasets.txt file using:
+```
+./download_datasets.sh
+```
+
+Then concatenate and resample the data into a single 16-kHz file:
+```
+./process_speech.sh
+```
+The script will produce an all_speech.pcm speech file in raw 16-bit PCM format.
+
+
+For data preparation you need to build Opus as detailed in the top-level README.
+You will need to use the --enable-dred configure option.
+The build will produce an executable named "dump_data".
+To prepare the training data, run:
+```
+./dump_data -train all_speech.pcm all_features.f32 /dev/null
+```
+
+## Training
+
+To perform training, run the following command:
+```
+python ./train_rdovae.py --sequence-length 400 --split-mode random_split --state-dim 80 --batch-size 512 --epochs 400 --lambda-max 0.04 --lr 0.003 --lr-decay-factor 0.0001 all_features.f32 output_dir
+```
+The final model will be in output_dir/checkpoints/chechpoint_400.pth.
+
+The model can be converted to C using:
+```
+python export_rdovae_weights.py output_dir/checkpoints/chechpoint_400.pth dred_c_dir
+```
+which will create a number of C source and header files in the fargan_c_dir directory.
+Copy these files to the opus/dnn/ directory (replacing the existing ones) and recompile Opus.
+
+## Inference
+
+DRED is integrated within the Opus codec and can be evaluated using the opus_demo
+executable. For example:
+```
+./opus_demo voip 16000 1 64000 -loss 50 -dred 100 -sim_loss 50 input.pcm output.pcm
+```
+Will tell the encoder to encode a 16 kHz raw audio file at 64 kb/s using up to 1 second
+of redundancy (units are based on 10-ms) and then simulate 50% loss. Refer to `opus_demo --help`
+for more details.

package/opus/dnn/torch/rdovae/download_datasets.sh

@@ -0,0 +1,6 @@
+mkdir datasets
+cd datasets
+for i in `grep https ../../../datasets.txt`
+do
+	wget $i
+done

package/opus/dnn/torch/rdovae/export_rdovae_weights.py

@@ -0,0 +1,365 @@
+"""
+/* Copyright (c) 2022 Amazon
+   Written by Jan Buethe */
+/*
+   Redistribution and use in source and binary forms, with or without
+   modification, are permitted provided that the following conditions
+   are met:
+
+   - Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+   - Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+
+   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
+   OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+"""
+
+import os
+import argparse
+import sys
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '../weight-exchange'))
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('checkpoint', type=str, help='rdovae model checkpoint')
+parser.add_argument('output_dir', type=str, help='output folder')
+parser.add_argument('--format', choices=['C', 'numpy'], help='output format, default: C', default='C')
+
+args = parser.parse_args()
+
+import torch
+import numpy as np
+
+from rdovae import RDOVAE
+from wexchange.torch import dump_torch_weights
+from wexchange.c_export import CWriter, print_vector
+
+def print_xml(xmlout, val, param, anchor, name):
+    xmlout.write(
+f"""
+            <table anchor="{anchor}_{name}">
+                <name>{param} values for {name}</name>
+                <thead>
+                    <tr><th>k</th><th>Q0</th><th>Q1</th><th>Q2</th><th>Q3</th><th>Q4</th><th>Q5</th><th>Q6</th><th>Q7</th><th>Q8</th><th>Q9</th><th>Q10</th><th>Q11</th><th>Q12</th><th>Q13</th><th>Q14</th><th>Q15</th></tr>
+                </thead>
+                <tbody>
+""")
+    for k in range(val.shape[1]):
+        xmlout.write(f"        <tr><th>{k}</th>")
+        for j in range(val.shape[0]):
+            xmlout.write(f"<th>{val[j][k]}</th>")
+        xmlout.write("</tr>\n")
+    xmlout.write(
+f"""
+                </tbody>
+            </table>
+""")
+def dump_statistical_model(writer, w, name, xmlout):
+    levels = w.shape[0]
+
+    print("printing statistical model")
+    quant_scales    = torch.nn.functional.softplus(w[:, 0, :]).numpy()
+    dead_zone       = 0.05 * torch.nn.functional.softplus(w[:, 1, :]).numpy()
+    r               = torch.sigmoid(w[:, 5 , :]).numpy()
+    p0              = torch.sigmoid(w[:, 4 , :]).numpy()
+    p0              = 1 - r ** (0.5 + 0.5 * p0)
+
+    scales_norm = 255./256./(1e-15+np.max(quant_scales,axis=0))
+    quant_scales = quant_scales*scales_norm
+    quant_scales_q8 = np.round(quant_scales * 2**8).astype(np.uint16)
+    dead_zone_q8   = np.clip(np.round(dead_zone * 2**8), 0, 255).astype(np.uint16)
+    r_q8           = np.clip(np.round(r * 2**8), 0, 255).astype(np.uint8)
+    p0_q8          = np.clip(np.round(p0 * 2**8), 0, 255).astype(np.uint16)
+
+    mask = (np.max(r_q8,axis=0) > 0) * (np.min(p0_q8,axis=0) < 255)
+    quant_scales_q8 = quant_scales_q8[:, mask]
+    dead_zone_q8 = dead_zone_q8[:, mask]
+    r_q8 = r_q8[:, mask]
+    p0_q8 = p0_q8[:, mask]
+    N = r_q8.shape[-1]
+
+    print_vector(writer.source, quant_scales_q8, f'dred_{name}_quant_scales_q8', dtype='opus_uint8', static=False)
+    print_vector(writer.source, dead_zone_q8, f'dred_{name}_dead_zone_q8', dtype='opus_uint8', static=False)
+    print_vector(writer.source, r_q8, f'dred_{name}_r_q8', dtype='opus_uint8', static=False)
+    print_vector(writer.source, p0_q8, f'dred_{name}_p0_q8', dtype='opus_uint8', static=False)
+
+    print_xml(xmlout, quant_scales_q8, "Scale", "scale", name)
+    print_xml(xmlout, dead_zone_q8, "Dead zone", "deadzone", name)
+    print_xml(xmlout, r_q8, "Decay (r)", "decay", name)
+    print_xml(xmlout, p0_q8, "P(0)", "p0", name)
+
+    writer.header.write(
+f"""
+extern const opus_uint8 dred_{name}_quant_scales_q8[{levels * N}];
+extern const opus_uint8 dred_{name}_dead_zone_q8[{levels * N}];
+extern const opus_uint8 dred_{name}_r_q8[{levels * N}];
+extern const opus_uint8 dred_{name}_p0_q8[{levels * N}];
+
+"""
+    )
+    return N, mask, torch.tensor(scales_norm[mask])
+
+
+def c_export(args, model):
+
+    message = f"Auto generated from checkpoint {os.path.basename(args.checkpoint)}"
+
+    enc_writer = CWriter(os.path.join(args.output_dir, "dred_rdovae_enc_data"), message=message, model_struct_name='RDOVAEEnc')
+    dec_writer = CWriter(os.path.join(args.output_dir, "dred_rdovae_dec_data"), message=message, model_struct_name='RDOVAEDec')
+    stats_writer = CWriter(os.path.join(args.output_dir, "dred_rdovae_stats_data"), message=message, enable_binary_blob=False)
+    constants_writer = CWriter(os.path.join(args.output_dir, "dred_rdovae_constants"), message=message, header_only=True, enable_binary_blob=False)
+    xmlout = open("stats.xml", "w")
+
+    # some custom includes
+    for writer in [enc_writer, dec_writer]:
+        writer.header.write(
+f"""
+#include "opus_types.h"
+
+#include "dred_rdovae.h"
+
+#include "dred_rdovae_constants.h"
+
+"""
+        )
+
+    stats_writer.header.write(
+f"""
+#include "opus_types.h"
+
+#include "dred_rdovae_constants.h"
+
+"""
+        )
+
+    latent_out = model.get_submodule('core_encoder.module.z_dense')
+    state_out = model.get_submodule('core_encoder.module.state_dense_2')
+    orig_latent_dim = latent_out.weight.shape[0]
+    orig_state_dim = state_out.weight.shape[0]
+    # statistical model
+    qembedding = model.statistical_model.quant_embedding.weight.detach()
+    levels = qembedding.shape[0]
+    qembedding = torch.reshape(qembedding, (levels, 6, -1))
+
+    latent_dim, latent_mask, latent_scale = dump_statistical_model(stats_writer, qembedding[:, :, :orig_latent_dim], 'latent', xmlout)
+    state_dim, state_mask, state_scale = dump_statistical_model(stats_writer, qembedding[:, :, orig_latent_dim:], 'state', xmlout)
+
+    padded_latent_dim = (latent_dim+7)//8*8
+    latent_pad = padded_latent_dim - latent_dim;
+    w = latent_out.weight[latent_mask,:]
+    w = w/latent_scale[:, None]
+    w = torch.cat([w, torch.zeros(latent_pad, w.shape[1])], dim=0)
+    b = latent_out.bias[latent_mask]
+    b = b/latent_scale
+    b = torch.cat([b, torch.zeros(latent_pad)], dim=0)
+    latent_out.weight = torch.nn.Parameter(w)
+    latent_out.bias = torch.nn.Parameter(b)
+
+    padded_state_dim = (state_dim+7)//8*8
+    state_pad = padded_state_dim - state_dim;
+    w = state_out.weight[state_mask,:]
+    w = w/state_scale[:, None]
+    w = torch.cat([w, torch.zeros(state_pad, w.shape[1])], dim=0)
+    b = state_out.bias[state_mask]
+    b = b/state_scale
+    b = torch.cat([b, torch.zeros(state_pad)], dim=0)
+    state_out.weight = torch.nn.Parameter(w)
+    state_out.bias = torch.nn.Parameter(b)
+
+    latent_in = model.get_submodule('core_decoder.module.dense_1')
+    state_in = model.get_submodule('core_decoder.module.hidden_init')
+    latent_in.weight = torch.nn.Parameter(latent_in.weight[:,latent_mask]*latent_scale)
+    state_in.weight = torch.nn.Parameter(state_in.weight[:,state_mask]*state_scale)
+
+    # encoder
+    encoder_dense_layers = [
+        ('core_encoder.module.dense_1'       , 'enc_dense1',   '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, _, quantize in encoder_dense_layers:
+        layer = model.get_submodule(name)
+        dump_torch_weights(enc_writer, layer, name=export_name, verbose=True, quantize=quantize, scale=None)
+
+
+    encoder_gru_layers = [
+        ('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=quantize, scale=None, recurrent_scale=None)
+                             for name, export_name, _, quantize in encoder_gru_layers])
+
+
+    encoder_conv_layers = [
+        ('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=quantize, scale=None) for name, export_name, _, quantize in encoder_conv_layers])
+
+
+    del enc_writer
+
+    # decoder
+    decoder_dense_layers = [
+        ('core_decoder.module.dense_1'      , 'dec_dense1',  'TANH', False),
+        ('core_decoder.module.glu1.gate'    , 'dec_glu1',    'TANH', True),
+        ('core_decoder.module.glu2.gate'    , 'dec_glu2',    'TANH', True),
+        ('core_decoder.module.glu3.gate'    , 'dec_glu3',    'TANH', True),
+        ('core_decoder.module.glu4.gate'    , 'dec_glu4',    'TANH', True),
+        ('core_decoder.module.glu5.gate'    , 'dec_glu5',    'TANH', True),
+        ('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', True),
+    ]
+
+    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, quantize=quantize, scale=None)
+
+
+    decoder_gru_layers = [
+        ('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=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),
+        ('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=quantize, scale=None) for name, export_name, _, quantize in decoder_conv_layers])
+
+    del dec_writer
+
+    del stats_writer
+
+    # constants
+    constants_writer.header.write(
+f"""
+#define DRED_NUM_FEATURES {model.feature_dim}
+
+#define DRED_LATENT_DIM {latent_dim}
+
+#define DRED_STATE_DIM {state_dim}
+
+#define DRED_PADDED_LATENT_DIM {padded_latent_dim}
+
+#define DRED_PADDED_STATE_DIM {padded_state_dim}
+
+#define DRED_NUM_QUANTIZATION_LEVELS {model.quant_levels}
+
+#define DRED_MAX_RNN_NEURONS {max(enc_max_rnn_units, dec_max_rnn_units)}
+
+#define DRED_MAX_CONV_INPUTS {max(enc_max_conv_inputs, dec_max_conv_inputs)}
+
+#define DRED_ENC_MAX_RNN_NEURONS {enc_max_conv_inputs}
+
+#define DRED_ENC_MAX_CONV_INPUTS {enc_max_conv_inputs}
+
+#define DRED_DEC_MAX_RNN_NEURONS {dec_max_rnn_units}
+
+"""
+    )
+
+    del constants_writer
+
+
+def numpy_export(args, model):
+
+    exchange_name_to_name = {
+        'encoder_stack_layer1_dense'    : 'core_encoder.module.dense_1',
+        'encoder_stack_layer3_dense'    : 'core_encoder.module.dense_2',
+        'encoder_stack_layer5_dense'    : 'core_encoder.module.dense_3',
+        'encoder_stack_layer7_dense'    : 'core_encoder.module.dense_4',
+        'encoder_stack_layer8_dense'    : 'core_encoder.module.dense_5',
+        'encoder_state_layer1_dense'    : 'core_encoder.module.state_dense_1',
+        'encoder_state_layer2_dense'    : 'core_encoder.module.state_dense_2',
+        'encoder_stack_layer2_gru'      : 'core_encoder.module.gru_1',
+        'encoder_stack_layer4_gru'      : 'core_encoder.module.gru_2',
+        'encoder_stack_layer6_gru'      : 'core_encoder.module.gru_3',
+        'encoder_stack_layer9_conv'     : 'core_encoder.module.conv1',
+        'statistical_model_embedding'   : 'statistical_model.quant_embedding',
+        'decoder_state1_dense'          : 'core_decoder.module.gru_1_init',
+        'decoder_state2_dense'          : 'core_decoder.module.gru_2_init',
+        'decoder_state3_dense'          : 'core_decoder.module.gru_3_init',
+        'decoder_stack_layer1_dense'    : 'core_decoder.module.dense_1',
+        'decoder_stack_layer3_dense'    : 'core_decoder.module.dense_2',
+        'decoder_stack_layer5_dense'    : 'core_decoder.module.dense_3',
+        'decoder_stack_layer7_dense'    : 'core_decoder.module.dense_4',
+        'decoder_stack_layer8_dense'    : 'core_decoder.module.dense_5',
+        'decoder_stack_layer9_dense'    : 'core_decoder.module.output',
+        'decoder_stack_layer2_gru'      : 'core_decoder.module.gru_1',
+        'decoder_stack_layer4_gru'      : 'core_decoder.module.gru_2',
+        'decoder_stack_layer6_gru'      : 'core_decoder.module.gru_3'
+    }
+
+    name_to_exchange_name = {value : key for key, value in exchange_name_to_name.items()}
+
+    for name, exchange_name in name_to_exchange_name.items():
+        print(f"printing layer {name}...")
+        dump_torch_weights(os.path.join(args.output_dir, exchange_name), model.get_submodule(name))
+
+
+if __name__ == "__main__":
+
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+
+    # load model from checkpoint
+    checkpoint = torch.load(args.checkpoint, map_location='cpu')
+    model = RDOVAE(*checkpoint['model_args'], **checkpoint['model_kwargs'])
+    missing_keys, unmatched_keys = model.load_state_dict(checkpoint['state_dict'], strict=False)
+    def _remove_weight_norm(m):
+        try:
+            torch.nn.utils.remove_weight_norm(m)
+        except ValueError:  # this module didn't have weight norm
+            return
+    model.apply(_remove_weight_norm)
+
+
+    if len(missing_keys) > 0:
+        raise ValueError(f"error: missing keys in state dict")
+
+    if len(unmatched_keys) > 0:
+        print(f"warning: the following keys were unmatched {unmatched_keys}")
+
+    if args.format == 'C':
+        c_export(args, model)
+    elif args.format == 'numpy':
+        numpy_export(args, model)
+    else:
+        raise ValueError(f'error: unknown export format {args.format}')