train_stateful_lstm_tf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d80fc465400cc842",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:17.956649100Z",
     "start_time": "2024-03-08T14:43:15.994120800Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import onnx\n",
    "import tf2onnx\n",
    "from scipy.io import wavfile\n",
    "from scipy.io.wavfile import write"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f57b51c451debb",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# 1. Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1cdc8701863f502b",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:19.241577300Z",
     "start_time": "2024-03-08T14:43:19.235040200Z"
    }
   },
   "outputs": [],
   "source": [
    "def prepare_training_data(config):\n",
    "    in_rate, in_data = wavfile.read(config[\"input_audio_path\"])\n",
    "    out_rate, out_data = wavfile.read(config[\"target_audio_path\"])\n",
    "    \n",
    "    if len(in_data) != len(out_data):\n",
    "        print(\"input and target files have different lengths\")\n",
    "        sys.exit()\n",
    "      \n",
    "    if len(in_data.shape) > 1 or len(out_data.shape) > 1:\n",
    "        print(\"expected mono files\")\n",
    "        sys.exit()\n",
    "\n",
    "    # Convert PCM16 to FP32\n",
    "    if in_data.dtype == \"int16\":\n",
    "        in_data = in_data / 32767\n",
    "        print(\"In data converted from PCM16 to FP32\")\n",
    "    if out_data.dtype == \"int16\":\n",
    "        out_data = out_data / 32767\n",
    "        print(\"Out data converted from PCM16 to FP32\")    \n",
    "\n",
    "    clean_data = in_data.astype(np.float32).flatten()\n",
    "    target_data = out_data.astype(np.float32).flatten()\n",
    "\n",
    "    # Split the data on a twenty percent mod\n",
    "    in_train, out_train, in_val, out_val = slice_on_mod(clean_data, target_data)\n",
    "\n",
    "    save_wav(config[\"output_path\"] + \"/train/\" + config[\"name\"] + \"-input.wav\", in_train)\n",
    "    save_wav(config[\"output_path\"] + \"/train/\" + config[\"name\"] + \"-target.wav\", out_train)\n",
    "\n",
    "    save_wav(config[\"output_path\"] + \"/test/\" + config[\"name\"] + \"-input.wav\", in_val)\n",
    "    save_wav(config[\"output_path\"] + \"/test/\" + config[\"name\"] + \"-target.wav\", out_val)\n",
    "\n",
    "    save_wav(config[\"output_path\"] + \"/val/\" + config[\"name\"] + \"-input.wav\", in_val)\n",
    "    save_wav(config[\"output_path\"] + \"/val/\" + config[\"name\"] + \"-target.wav\", out_val)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d2db02127a930a4c",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:20.837448800Z",
     "start_time": "2024-03-08T14:43:20.814913400Z"
    }
   },
   "outputs": [],
   "source": [
    "def slice_on_mod(input_data, target_data, mod=5):\n",
    "    # Split the data on a modulus.\n",
    "\n",
    "    # Type cast to an integer the modulus\n",
    "    mod = int(mod)\n",
    "\n",
    "    # Split the data into 100 pieces\n",
    "    input_split = np.array_split(input_data, 100)\n",
    "    target_split = np.array_split(target_data, 100)\n",
    "\n",
    "    val_input_data = []\n",
    "    val_target_data = []\n",
    "    # Traverse the range of the indexes of the input signal reversed and pop every 5th for val\n",
    "    for i in reversed(range(len(input_split))):\n",
    "        if i % mod == 0:\n",
    "            # Store the validation data\n",
    "            val_input_data.append(input_split[i])\n",
    "            val_target_data.append(target_split[i])\n",
    "            # Remove the validation data from training\n",
    "            input_split.pop(i)\n",
    "            target_split.pop(i)\n",
    "\n",
    "    # Flatten val_data down to one dimension and concatenate\n",
    "    val_input_data = np.concatenate(val_input_data)\n",
    "    val_target_data = np.concatenate(val_target_data)\n",
    "\n",
    "    # Concatenate back together\n",
    "    training_input_data = np.concatenate(input_split)\n",
    "    training_target_data = np.concatenate(target_split)\n",
    "    return training_input_data, training_target_data, val_input_data, val_target_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ac5a30c5ce341b1",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:22.620429600Z",
     "start_time": "2024-03-08T14:43:22.599885400Z"
    }
   },
   "outputs": [],
   "source": [
    "def save_wav(name, data):\n",
    "    directory = os.path.dirname(name)\n",
    "\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "        \n",
    "    wavfile.write(name, 44100, data.flatten().astype(np.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c7cc0b6f364b5692",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:23.547006800Z",
     "start_time": "2024-03-08T14:43:23.400487100Z"
    }
   },
   "outputs": [],
   "source": [
    "importConfig = {\n",
    "    \"input_audio_path\": \"TrainingData/ts9-input.wav\",\n",
    "    \"target_audio_path\": \"TrainingData/ts9-target.wav\",\n",
    "    \"output_path\": \"Data\",\n",
    "    \"name\": \"ts9\"\n",
    "}\n",
    "\n",
    "prepare_training_data(importConfig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46db2dbe23cb7b1c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c026cfbd4fba6c7c",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:24.872267400Z",
     "start_time": "2024-03-08T14:43:24.843712Z"
    }
   },
   "outputs": [],
   "source": [
    "class DataSet:\n",
    "    def __init__(self, data_dir='Data/'):\n",
    "        self.data_dir = data_dir\n",
    "        self.subsets = {}\n",
    "\n",
    "    def create_subset(self, name, frame_len=22050):\n",
    "        self.subsets[name] = {'input': None, 'target': None, 'frame_len': frame_len}\n",
    "\n",
    "    def load_file(self, subset_name, base_filename):\n",
    "        if subset_name not in self.subsets:\n",
    "            raise ValueError(f\"Subset '{subset_name}' does not exist\")\n",
    "\n",
    "        input_file = os.path.join(self.data_dir, f\"{base_filename}-input.wav\")\n",
    "        target_file = os.path.join(self.data_dir, f\"{base_filename}-target.wav\")\n",
    "\n",
    "        try:\n",
    "            self.subsets[subset_name]['input'] = self.load_and_process(input_file, self.subsets[subset_name]['frame_len'])\n",
    "            self.subsets[subset_name]['target'] = self.load_and_process(target_file, self.subsets[subset_name]['frame_len'])\n",
    "        except FileNotFoundError as e:\n",
    "            print(f\"File Not Found: {e.filename}\")\n",
    "            return\n",
    "\n",
    "    def load_and_process(self, file_path, frame_len):\n",
    "        sample_rate, data = wavfile.read(file_path)\n",
    "        data = data.astype(np.float32)\n",
    "        return self.framify(data, frame_len)\n",
    "\n",
    "    def framify(self, audio, frame_len):\n",
    "        seg_num = math.ceil(audio.shape[0] / frame_len)\n",
    "        padded_length = seg_num * frame_len\n",
    "        padded_audio = np.pad(audio, (0, padded_length - audio.shape[0]), mode='constant')\n",
    "\n",
    "        reshaped_audio = np.reshape(padded_audio, (seg_num, frame_len, 1))\n",
    "        return tf.convert_to_tensor(reshaped_audio, dtype=tf.float32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61f2ed6f6e582365",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "860763dc7078f73d",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:26.604768400Z",
     "start_time": "2024-03-08T14:43:26.570654200Z"
    }
   },
   "outputs": [],
   "source": [
    "class StatefulLSTM(tf.keras.Model):\n",
    "    def __init__(self, input_size=1, output_size=1, hidden_size=32, skip=1, bias_fl=True, batch_size=4096):\n",
    "        super(StatefulLSTM, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.skip = skip\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        self.lstm = tf.keras.layers.LSTM(units=hidden_size, return_sequences=True, stateful=True, return_state=True, batch_size=batch_size, kernel_initializer='glorot_uniform', recurrent_initializer='zeros', bias_initializer='zeros', use_bias=bias_fl)    \n",
    "        self.dense = tf.keras.layers.Dense(units=output_size, activation=None, batch_size=batch_size, use_bias=bias_fl, kernel_initializer='glorot_uniform', bias_initializer='zeros')\n",
    "        \n",
    "        # Build LSTM before training, because stateful lstm requires information batch size to build static graph\n",
    "        self.lstm.build((batch_size, input_size, 1))\n",
    "        # Set the biases to zero, seems to be not working with bias_initializer='zeros'\n",
    "        self.lstm.weights[2].assign(tf.zeros((hidden_size * 4)))\n",
    "        \n",
    "    def call(self, x):\n",
    "        if self.skip:\n",
    "            # save the residual for the skip connection\n",
    "            res = x[:, :, 0:self.skip]\n",
    "            x, _, _ = self.lstm(x)\n",
    "            return self.dense(x) + res\n",
    "        else:\n",
    "            x, _, _ = self.lstm(x)\n",
    "            x = self.dense(x)\n",
    "            return x\n",
    "            \n",
    "    def reset_hidden(self, batch_size):\n",
    "        initial_state = [tf.zeros((self.batch_size ,self.hidden_size), dtype=tf.float32), tf.zeros((self.batch_size, self.hidden_size), dtype=tf.float32)]\n",
    "        self.lstm.reset_states(states=initial_state)\n",
    "    \n",
    "    def train_epoch(self, input_data, target_data, loss_fcn, optim, bs, init_len=200, up_fr=1000):\n",
    "\n",
    "        # shuffle the segments at the start of the epoch\n",
    "        shuffle = tf.random.shuffle(tf.range(input_data.shape[0]))\n",
    "    \n",
    "        self.reset_hidden(bs)\n",
    "\n",
    "        # Iterate over the batches\n",
    "        ep_loss = 0\n",
    "        for batch_i in range(math.ceil(shuffle.shape[0] / bs)):\n",
    "            if shuffle[batch_i * bs:(batch_i + 1) * bs].shape[0] != bs:\n",
    "                # If the final batch is smaller than the batch size, break the loop\n",
    "                break \n",
    "\n",
    "            # Use tf.gather to index the tensors\n",
    "            input_batch = tf.gather(input_data, shuffle[batch_i * bs:(batch_i + 1) * bs], axis=0)\n",
    "            target_batch = tf.gather(target_data, shuffle[batch_i * bs:(batch_i + 1) * bs], axis=0)\n",
    "            \n",
    "            # Initialise network hidden state by processing some samples then zero the gradient buffers\n",
    "            # For training processing eine Anfangssequenz, damit ein brauchbarer hidden state vorliegt\n",
    "            # Training startet erst nach! einem eingelaufen hidden state\n",
    "            self(input_batch[:, 0:init_len, :])\n",
    "        \n",
    "            start_i = init_len\n",
    "            batch_loss = 0\n",
    "            # Iterate over the remaining samples in the mini batch\n",
    "            for k in range(math.ceil((input_batch.shape[1] - init_len) / up_fr)):\n",
    "                \n",
    "                with tf.GradientTape() as g:\n",
    "                    # Process input batch with neural network    \n",
    "                    output = self(input_batch[:, start_i:start_i + up_fr, :])\n",
    "                    loss = loss_fcn(output, target_batch[:, start_i:start_i + up_fr, :])\n",
    "                    with g.stop_recording():\n",
    "                        dloss_dw = g.gradient(loss, self.trainable_variables)\n",
    "                        optim.apply_gradients(zip(dloss_dw, self.trainable_variables))\n",
    "                    g.reset()\n",
    "                        \n",
    "                print(f\"loss: {loss}\")\n",
    "\n",
    "                # Update the start index for the next iteration and add the loss to the batch_loss total\n",
    "                start_i += up_fr\n",
    "                batch_loss += loss\n",
    "\n",
    "            # Add the average batch loss to the epoch loss and reset the hidden states to zeros\n",
    "            ep_loss += batch_loss / (k + 1)\n",
    "            self.reset_hidden(bs)\n",
    "        \n",
    "        return ep_loss / (batch_i + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cbd60e7e671135d4",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:29.621084Z",
     "start_time": "2024-03-08T14:43:29.603532300Z"
    }
   },
   "outputs": [],
   "source": [
    "class ESRLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self):\n",
    "        super(ESRLoss, self).__init__()\n",
    "        self.epsilon = 0.00001\n",
    "\n",
    "    def call(self, output, target):\n",
    "        loss = tf.add(target, -output)\n",
    "        loss = tf.pow(loss, 2)\n",
    "        loss = tf.reduce_mean(loss)\n",
    "        energy = tf.reduce_mean(tf.pow(target, 2)) + self.epsilon\n",
    "        loss = tf.divide(loss, energy)\n",
    "        return loss\n",
    "    \n",
    "class DCLoss(tf.keras.losses.Loss):\n",
    "    def __init__(self):\n",
    "        super(DCLoss, self).__init__()\n",
    "        self.epsilon = 0.00001\n",
    "\n",
    "    def call(self, output, target):\n",
    "        loss = tf.pow(tf.add(tf.reduce_mean(target, axis=0), -tf.reduce_mean(output, axis=0)), 2)\n",
    "        loss = tf.reduce_mean(loss)\n",
    "        energy = tf.reduce_mean(tf.pow(target, 2)) + self.epsilon\n",
    "        loss = tf.divide(loss, energy)\n",
    "        return loss\n",
    "    \n",
    "class LossWrapper(tf.keras.losses.Loss):\n",
    "    def __init__(self, losses):\n",
    "        super(LossWrapper, self).__init__()\n",
    "        loss_dict = {'ESR': ESRLoss(), 'DC': DCLoss()}\n",
    "\n",
    "        loss_functions = [[loss_dict[key], value] for key, value in losses.items()]\n",
    "\n",
    "        self.loss_functions = tuple([items[0] for items in loss_functions])\n",
    "        try:\n",
    "            self.loss_factors = tuple(tf.constant([items[1] for items in loss_functions]))\n",
    "        except IndexError:\n",
    "            self.loss_factors = tf.ones(len(self.loss_functions))\n",
    "\n",
    "    def call(self, output, target):\n",
    "        loss = 0\n",
    "        for i, losses in enumerate(self.loss_functions):\n",
    "            loss += tf.multiply(losses(output, target), self.loss_factors[i])\n",
    "        return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1c0c26066b7bc449",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:31.746171100Z",
     "start_time": "2024-03-08T14:43:31.736097100Z"
    }
   },
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"input_size\": 1, # Number of channels\n",
    "    \"output_size\": 1, # Number of channels\n",
    "    \"skip_con\": 1, # is there a skip connection for the input to the output\n",
    "    \"epochs\": 20,\n",
    "    \"batch_size\": 50,\n",
    "    \"init_length\": 200, # Number of sequence samples to process before starting weight updates\n",
    "    \"up_fr\": 1000, # For recurrent models, number of samples to run in between updating network weights\n",
    "    \"learning_rate\": 0.005,\n",
    "    \"hidden_size\": 20,\n",
    "    \"loss_fcns\": {\"ESR\": 0.75, \"DC\": 0.25},\n",
    "    \"hardware_device\": \"ts9\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cca470a4e37d92cd",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:33.181196500Z",
     "start_time": "2024-03-08T14:43:32.937713100Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "These are the physical devices available:\n",
      "[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]\n",
      "These are the visible devices:\n",
      "[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]\n",
      "Creating Stateful LSTM\n",
      "Model: \"stateful_lstm\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " lstm (LSTM)                 multiple                  1760      \n",
      "                                                                 \n",
      " dense (Dense)               multiple                  21        \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 1781 (6.96 KB)\n",
      "Trainable params: 1781 (6.96 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "physical_devices = tf.config.list_physical_devices()\n",
    "print(f\"These are the physical devices available:\\n{physical_devices}\")\n",
    "\n",
    "try:\n",
    "    # Disable all GPUS\n",
    "    tf.config.set_visible_devices([], 'GPU')\n",
    "    visible_devices = tf.config.get_visible_devices()\n",
    "    print(f\"These are the visible devices:\\n{visible_devices}\")\n",
    "except:\n",
    "    pass\n",
    "    \n",
    "print(\"Creating Stateful LSTM\")\n",
    "network = StatefulLSTM(input_size=config[\"input_size\"], \n",
    "                       output_size=config[\"output_size\"], \n",
    "                       hidden_size=config[\"hidden_size\"], \n",
    "                       skip=config[\"skip_con\"],\n",
    "                       batch_size=config[\"batch_size\"])\n",
    "\n",
    "optimiser = tf.keras.optimizers.Adam(learning_rate=config[\"learning_rate\"], weight_decay=1e-4, epsilon=1e-8)\n",
    "# loss_functions = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.SUM)\n",
    "\n",
    "loss_functions = LossWrapper(config[\"loss_fcns\"])\n",
    "\n",
    "network.compile(optimizer=optimiser, loss=loss_functions)\n",
    "network.build((config[\"batch_size\"],1,1))\n",
    "network.summary()        \n",
    "\n",
    "# Usage Example\n",
    "dataset = DataSet()\n",
    "dataset.create_subset('train', frame_len=22050)\n",
    "dataset.load_file('train', os.path.join('train', config[\"hardware_device\"]))\n",
    "\n",
    "dataset.create_subset('val')\n",
    "dataset.load_file('val', os.path.join('val', config[\"hardware_device\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959e152f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(1, config[\"epochs\"] + 1):\n",
    "    print(\"Epoch: \", epoch)\n",
    "\n",
    "    epoch_loss = network.train_epoch(dataset.subsets['train']['input'],\n",
    "                                     dataset.subsets['train']['target'],\n",
    "                                     loss_functions,\n",
    "                                     optimiser,\n",
    "                                     config['batch_size'],\n",
    "                                     config['init_length'],\n",
    "                                     config['up_fr'])\n",
    "\n",
    "    print(\"Epoch loss:\", epoch_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc03fcf5a2209c3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T23:12:32.907471700Z",
     "start_time": "2024-01-21T23:12:32.767844300Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "name = \"model_0\"\n",
    "network.save_weights(\"models/\"+name+\"/stateful-lstm\")\n",
    "network.save_weights('models/lstm_model.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "dc52b05d7df1ea3a",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:47:20.744178500Z",
     "start_time": "2024-03-08T14:47:20.603878600Z"
    }
   },
   "outputs": [],
   "source": [
    "inference_batch_size = 1\n",
    "\n",
    "name = \"model_0\"\n",
    "\n",
    "inference_model = StatefulLSTM(input_size=1, \n",
    "                               output_size=1,\n",
    "                               hidden_size=20,\n",
    "                               skip=0,\n",
    "                               batch_size=inference_batch_size)\n",
    "\n",
    "input_shape = (inference_batch_size,2048,1)\n",
    "\n",
    "inference_model.build(input_shape)\n",
    "\n",
    "inference_model.load_weights('models/lstm_model.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7d6cd479",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:49.356222200Z",
     "start_time": "2024-03-08T14:43:49.333548600Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(299, 22050, 1)\n",
      "(299, 22050, 1)\n"
     ]
    }
   ],
   "source": [
    "print(dataset.subsets['train']['input'].shape)\n",
    "input = dataset.subsets['train']['input']\n",
    "print(input.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "09c7d381",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-08T14:43:56.505202Z",
     "start_time": "2024-03-08T14:43:50.901351700Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 10/299 [>.............................] - ETA: 1:52"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[13], line 2\u001B[0m\n\u001B[0;32m      1\u001B[0m inference_model\u001B[38;5;241m.\u001B[39mreset_hidden(inference_batch_size)\n\u001B[1;32m----> 2\u001B[0m output \u001B[38;5;241m=\u001B[39m \u001B[43minference_model\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mpredict\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mbatch_size\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minference_batch_size\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\keras\\src\\utils\\traceback_utils.py:65\u001B[0m, in \u001B[0;36mfilter_traceback.<locals>.error_handler\u001B[1;34m(*args, **kwargs)\u001B[0m\n\u001B[0;32m     63\u001B[0m filtered_tb \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[0;32m     64\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m---> 65\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfn\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m     66\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[0;32m     67\u001B[0m     filtered_tb \u001B[38;5;241m=\u001B[39m _process_traceback_frames(e\u001B[38;5;241m.\u001B[39m__traceback__)\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\keras\\src\\engine\\training.py:2554\u001B[0m, in \u001B[0;36mModel.predict\u001B[1;34m(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)\u001B[0m\n\u001B[0;32m   2552\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m step \u001B[38;5;129;01min\u001B[39;00m data_handler\u001B[38;5;241m.\u001B[39msteps():\n\u001B[0;32m   2553\u001B[0m     callbacks\u001B[38;5;241m.\u001B[39mon_predict_batch_begin(step)\n\u001B[1;32m-> 2554\u001B[0m     tmp_batch_outputs \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mpredict_function\u001B[49m\u001B[43m(\u001B[49m\u001B[43miterator\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m   2555\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m data_handler\u001B[38;5;241m.\u001B[39mshould_sync:\n\u001B[0;32m   2556\u001B[0m         context\u001B[38;5;241m.\u001B[39masync_wait()\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\util\\traceback_utils.py:150\u001B[0m, in \u001B[0;36mfilter_traceback.<locals>.error_handler\u001B[1;34m(*args, **kwargs)\u001B[0m\n\u001B[0;32m    148\u001B[0m filtered_tb \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[0;32m    149\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m--> 150\u001B[0m   \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfn\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    151\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[0;32m    152\u001B[0m   filtered_tb \u001B[38;5;241m=\u001B[39m _process_traceback_frames(e\u001B[38;5;241m.\u001B[39m__traceback__)\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\polymorphic_function\\polymorphic_function.py:825\u001B[0m, in \u001B[0;36mFunction.__call__\u001B[1;34m(self, *args, **kwds)\u001B[0m\n\u001B[0;32m    822\u001B[0m compiler \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mxla\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_jit_compile \u001B[38;5;28;01melse\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mnonXla\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    824\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m OptionalXlaContext(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_jit_compile):\n\u001B[1;32m--> 825\u001B[0m   result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwds\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    827\u001B[0m new_tracing_count \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mexperimental_get_tracing_count()\n\u001B[0;32m    828\u001B[0m without_tracing \u001B[38;5;241m=\u001B[39m (tracing_count \u001B[38;5;241m==\u001B[39m new_tracing_count)\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\polymorphic_function\\polymorphic_function.py:864\u001B[0m, in \u001B[0;36mFunction._call\u001B[1;34m(self, *args, **kwds)\u001B[0m\n\u001B[0;32m    861\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_lock\u001B[38;5;241m.\u001B[39mrelease()\n\u001B[0;32m    862\u001B[0m \u001B[38;5;66;03m# In this case we have not created variables on the first call. So we can\u001B[39;00m\n\u001B[0;32m    863\u001B[0m \u001B[38;5;66;03m# run the first trace but we should fail if variables are created.\u001B[39;00m\n\u001B[1;32m--> 864\u001B[0m results \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_variable_creation_fn\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwds\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    865\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_created_variables \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m ALLOW_DYNAMIC_VARIABLE_CREATION:\n\u001B[0;32m    866\u001B[0m   \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mCreating variables on a non-first call to a function\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    867\u001B[0m                    \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m decorated with tf.function.\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\polymorphic_function\\tracing_compiler.py:148\u001B[0m, in \u001B[0;36mTracingCompiler.__call__\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m    145\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_lock:\n\u001B[0;32m    146\u001B[0m   (concrete_function,\n\u001B[0;32m    147\u001B[0m    filtered_flat_args) \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_maybe_define_function(args, kwargs)\n\u001B[1;32m--> 148\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mconcrete_function\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_flat\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m    149\u001B[0m \u001B[43m    \u001B[49m\u001B[43mfiltered_flat_args\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcaptured_inputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mconcrete_function\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcaptured_inputs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\polymorphic_function\\monomorphic_function.py:1349\u001B[0m, in \u001B[0;36mConcreteFunction._call_flat\u001B[1;34m(self, args, captured_inputs)\u001B[0m\n\u001B[0;32m   1345\u001B[0m possible_gradient_type \u001B[38;5;241m=\u001B[39m gradients_util\u001B[38;5;241m.\u001B[39mPossibleTapeGradientTypes(args)\n\u001B[0;32m   1346\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m (possible_gradient_type \u001B[38;5;241m==\u001B[39m gradients_util\u001B[38;5;241m.\u001B[39mPOSSIBLE_GRADIENT_TYPES_NONE\n\u001B[0;32m   1347\u001B[0m     \u001B[38;5;129;01mand\u001B[39;00m executing_eagerly):\n\u001B[0;32m   1348\u001B[0m   \u001B[38;5;66;03m# No tape is watching; skip to running the function.\u001B[39;00m\n\u001B[1;32m-> 1349\u001B[0m   \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_build_call_outputs(\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_inference_function\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[0;32m   1350\u001B[0m forward_backward \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_select_forward_and_backward_functions(\n\u001B[0;32m   1351\u001B[0m     args,\n\u001B[0;32m   1352\u001B[0m     possible_gradient_type,\n\u001B[0;32m   1353\u001B[0m     executing_eagerly)\n\u001B[0;32m   1354\u001B[0m forward_function, args_with_tangents \u001B[38;5;241m=\u001B[39m forward_backward\u001B[38;5;241m.\u001B[39mforward()\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\polymorphic_function\\atomic_function.py:196\u001B[0m, in \u001B[0;36mAtomicFunction.__call__\u001B[1;34m(self, *args)\u001B[0m\n\u001B[0;32m    194\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m record\u001B[38;5;241m.\u001B[39mstop_recording():\n\u001B[0;32m    195\u001B[0m   \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_bound_context\u001B[38;5;241m.\u001B[39mexecuting_eagerly():\n\u001B[1;32m--> 196\u001B[0m     outputs \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_bound_context\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcall_function\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m    197\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mname\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    198\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43mlist\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43margs\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    199\u001B[0m \u001B[43m        \u001B[49m\u001B[38;5;28;43mlen\u001B[39;49m\u001B[43m(\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mfunction_type\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mflat_outputs\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m    200\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m    201\u001B[0m   \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m    202\u001B[0m     outputs \u001B[38;5;241m=\u001B[39m make_call_op_in_graph(\u001B[38;5;28mself\u001B[39m, \u001B[38;5;28mlist\u001B[39m(args))\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\context.py:1457\u001B[0m, in \u001B[0;36mContext.call_function\u001B[1;34m(self, name, tensor_inputs, num_outputs)\u001B[0m\n\u001B[0;32m   1455\u001B[0m cancellation_context \u001B[38;5;241m=\u001B[39m cancellation\u001B[38;5;241m.\u001B[39mcontext()\n\u001B[0;32m   1456\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m cancellation_context \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[1;32m-> 1457\u001B[0m   outputs \u001B[38;5;241m=\u001B[39m \u001B[43mexecute\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mexecute\u001B[49m\u001B[43m(\u001B[49m\n\u001B[0;32m   1458\u001B[0m \u001B[43m      \u001B[49m\u001B[43mname\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdecode\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43mutf-8\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1459\u001B[0m \u001B[43m      \u001B[49m\u001B[43mnum_outputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mnum_outputs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1460\u001B[0m \u001B[43m      \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtensor_inputs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1461\u001B[0m \u001B[43m      \u001B[49m\u001B[43mattrs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mattrs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1462\u001B[0m \u001B[43m      \u001B[49m\u001B[43mctx\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\n\u001B[0;32m   1463\u001B[0m \u001B[43m  \u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m   1464\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m   1465\u001B[0m   outputs \u001B[38;5;241m=\u001B[39m execute\u001B[38;5;241m.\u001B[39mexecute_with_cancellation(\n\u001B[0;32m   1466\u001B[0m       name\u001B[38;5;241m.\u001B[39mdecode(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mutf-8\u001B[39m\u001B[38;5;124m\"\u001B[39m),\n\u001B[0;32m   1467\u001B[0m       num_outputs\u001B[38;5;241m=\u001B[39mnum_outputs,\n\u001B[1;32m   (...)\u001B[0m\n\u001B[0;32m   1471\u001B[0m       cancellation_manager\u001B[38;5;241m=\u001B[39mcancellation_context,\n\u001B[0;32m   1472\u001B[0m   )\n",
      "File \u001B[1;32mD:\\Code\\projects\\python\\stateful-lstm\\venv\\lib\\site-packages\\tensorflow\\python\\eager\\execute.py:53\u001B[0m, in \u001B[0;36mquick_execute\u001B[1;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001B[0m\n\u001B[0;32m     51\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m     52\u001B[0m   ctx\u001B[38;5;241m.\u001B[39mensure_initialized()\n\u001B[1;32m---> 53\u001B[0m   tensors \u001B[38;5;241m=\u001B[39m \u001B[43mpywrap_tfe\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mTFE_Py_Execute\u001B[49m\u001B[43m(\u001B[49m\u001B[43mctx\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_handle\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdevice_name\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mop_name\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m     54\u001B[0m \u001B[43m                                      \u001B[49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mattrs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mnum_outputs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m     55\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m core\u001B[38;5;241m.\u001B[39m_NotOkStatusException \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[0;32m     56\u001B[0m   \u001B[38;5;28;01mif\u001B[39;00m name \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "inference_model.reset_hidden(inference_batch_size)\n",
    "output = inference_model.predict(input, batch_size=inference_batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d9b66c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_concat = output[0, :, :]\n",
    "input_concat = input[0, :, :]\n",
    "for i in range(1, output.shape[0]):\n",
    "    output_concat =  np.concatenate((output_concat, output[i, :, :]), axis=0)\n",
    "    input_concat = np.concatenate((input_concat, input[i, :, :]), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd87bfc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "write(os.path.join(\"input-tensorflow.wav\"), 44100, input_concat.reshape(-1, 1))\n",
    "write(os.path.join(\"output-tensorflow.wav\"), 44100, output_concat.reshape(-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e49f7edcd3489db9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T18:01:01.116143400Z",
     "start_time": "2024-01-21T18:01:00.163057300Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "input_shape = (1,1,1)\n",
    "test_sequence = tf.zeros(input_shape)\n",
    "\n",
    "\n",
    "print(\"Running prediction..\")\n",
    "prediction = inference_model.predict(test_sequence)\n",
    "print(f\"prediction {prediction}\")\n",
    "\n",
    "print(\"Running prediction..\")\n",
    "prediction = inference_model.predict(test_sequence)\n",
    "print(f\"prediction2 {prediction}\")\n",
    "\n",
    "print(\"test_sequence shape: \", test_sequence.shape)\n",
    "print(\"prediction_2 shape: \", prediction.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8195b6ac60138cd3",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-08T14:47:30.312518500Z",
     "start_time": "2024-03-08T14:47:26.006963500Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: C:\\Users\\VALENT~1\\AppData\\Local\\Temp\\tmp4c41jl5e\\assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: C:\\Users\\VALENT~1\\AppData\\Local\\Temp\\tmp4c41jl5e\\assets\n"
     ]
    }
   ],
   "source": [
    "inference_model.reset_hidden(inference_batch_size)\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "func = tf.function(inference_model).get_concrete_function(\n",
    "    tf.TensorSpec(input_shape, dtype=tf.float32))\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([func], inference_model)\n",
    "tflite_model = converter.convert()\n",
    "\n",
    "# Save the model.\n",
    "with open(\"models/\"+name+\"/stateful-lstm-dynamic.tflite\", 'wb') as f:\n",
    "  f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tf2onnx.tf_loader:TF freezing failed. Attempting to fix freezing errors.\n",
      "WARNING:tf2onnx.tf_loader:Removed Mul stateful_lstm_1/lstm_1/AssignVariableOp_1\n",
      "WARNING:tf2onnx.tf_loader:Removed Mul stateful_lstm_1/lstm_1/AssignVariableOp\n"
     ]
    }
   ],
   "source": [
    "input_shape = [1, None, 1]\n",
    "\n",
    "# Define the input shape\n",
    "input_signature = [tf.TensorSpec(input_shape, tf.float32, name='input')]\n",
    "\n",
    "# Convert the model\n",
    "onnx_model, _ = tf2onnx.convert.from_keras(inference_model, input_signature, opset=13)\n",
    "onnx.save(proto=onnx_model, f=\"models/\"+name+\"/\"+\"stateful-lstm-tflite-dynamic.onnx\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "48adc781"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78c5b579",
   "metadata": {},
   "outputs": [],
   "source": [
    "tflite_model_path = \"models/\"+name+\"/stateful-lstm.tflite\"\n",
    "interpreter = tf.lite.Interpreter(model_path=tflite_model_path)\n",
    "interpreter.allocate_tensors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc095210",
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.set_tensor(interpreter.get_input_details()[0][\"index\"], tf.zeros((1, 2048, 1)))\n",
    "interpreter.invoke()\n",
    "prediction = interpreter.get_tensor(interpreter.get_output_details()[0][\"index\"])\n",
    "print(f\"prediction3 {prediction}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1994958",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "import numpy as np\n",
    "\n",
    "# Load the ONNX model\n",
    "onnx_model_path = \"models/\"+name+\"/\"+\"stateful-lstm-tflite.onnx\"\n",
    "ort_session = onnxruntime.InferenceSession(onnx_model_path)\n",
    "\n",
    "# Create an input tensor with shape (1, 1, 1) filled with zeros\n",
    "test_input = np.zeros((1, 2048, 1), dtype=np.float32)\n",
    "\n",
    "# Run inference on the ONNX model\n",
    "ort_inputs = {\"input\": test_input}\n",
    "ort_outputs = ort_session.run(None, ort_inputs)\n",
    "ort_outputs2 = ort_session.run(None, ort_inputs)\n",
    "\n",
    "# Print the output\n",
    "print(ort_outputs)\n",
    "# Print the output\n",
    "print(ort_outputs2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7efe83d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an input tensor with shape (1, 1, 1) filled with zeros\n",
    "test_input = np.zeros((1, 2048, 1), dtype=np.float32)\n",
    "\n",
    "# Run inference on the ONNX model\n",
    "ort_inputs = {\"input\": test_input}\n",
    "ort_outputs = ort_session.run(None, ort_inputs)\n",
    "ort_outputs2 = ort_session.run(None, ort_inputs)\n",
    "# Print the output\n",
    "print(ort_outputs)\n",
    "# Print the output\n",
    "print(ort_outputs2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}