@scrypted/libav 1.0.169 → 1.0.171

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

@@ -1,145 +0,0 @@
-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

@@ -1,55 +0,0 @@
-# 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

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

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

@@ -1,365 +0,0 @@
-"""
-/* 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}')