train_stateful_lstm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d80fc465400cc842",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:31.854186300Z",
     "start_time": "2024-01-21T13:48:26.741021400Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import math\n",
    "import torch\n",
    "import torchinfo\n",
    "import torch.nn as nn\n",
    "import onnx\n",
    "import numpy as np\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": null,
   "id": "1cdc8701863f502b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:31.871814600Z",
     "start_time": "2024-01-21T13:48:31.860178600Z"
    },
    "collapsed": false
   },
   "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": null,
   "id": "d2db02127a930a4c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:31.899871400Z",
     "start_time": "2024-01-21T13:48:31.876832600Z"
    },
    "collapsed": false
   },
   "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": null,
   "id": "8ac5a30c5ce341b1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:31.904386100Z",
     "start_time": "2024-01-21T13:48:31.888337300Z"
    },
    "collapsed": false
   },
   "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": null,
   "id": "c7cc0b6f364b5692",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:32.399235300Z",
     "start_time": "2024-01-21T13:48:31.905387800Z"
    },
    "collapsed": false
   },
   "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": "23b0f24e4fb42f7d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "778eadd3a8b28ea4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:53:27.140718600Z",
     "start_time": "2024-01-21T13:53:27.084422200Z"
    },
    "collapsed": false
   },
   "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 framify(self, audio, frame_len):\n",
    "        audio = np.expand_dims(audio, 1) if len(audio.shape) == 1 else audio\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]), (0, 0)), mode='constant')\n",
    "\n",
    "        reshaped_audio = np.reshape(padded_audio, (seg_num, frame_len, audio.shape[1]))\n",
    "        return torch.from_numpy(reshaped_audio).float()\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61f2ed6f6e582365",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "860763dc7078f73d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T14:34:08.256498400Z",
     "start_time": "2024-01-21T14:34:08.221258200Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class StatefulLSTM(nn.Module):\n",
    "    def __init__(self, input_size=1, output_size=1, hidden_size=32, skip=0, bias_fl=True):\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",
    "\n",
    "        # Create dictionary of possible block types\n",
    "        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size)\n",
    "        self.lin = nn.Linear(hidden_size, output_size, bias=bias_fl)\n",
    "        self.hidden = (torch.zeros(self.input_size, 1, self.hidden_size),\n",
    "                       torch.zeros(self.input_size, 1, self.hidden_size))\n",
    "        torch.nn.init.xavier_uniform_(self.lstm.weight_hh_l0)\n",
    "        torch.nn.init.xavier_uniform_(self.lstm.weight_ih_l0)\n",
    "        torch.nn.init.zeros_(self.lstm.bias_hh_l0)\n",
    "        torch.nn.init.zeros_(self.lstm.bias_ih_l0)\n",
    "        torch.nn.init.xavier_uniform_(self.lin.weight)\n",
    "        torch.nn.init.zeros_(self.lin.bias)\n",
    "\n",
    "    # Origin forward function\n",
    "    def forward(self, x):\n",
    "        if self.skip:\n",
    "            # save the residual for the skip connection\n",
    "            res = x[:, :, 0:self.skip]\n",
    "            x, self.hidden = self.lstm(x, self.hidden)\n",
    "            return self.lin(x) + res\n",
    "        else:\n",
    "            x, self.hidden = self.lstm(x, self.hidden)\n",
    "            x = self.lin(x)\n",
    "            return x\n",
    "\n",
    "    # detach hidden state, this resets gradient tracking on the hidden state\n",
    "    def detach_hidden(self):\n",
    "        if self.hidden.__class__ == tuple:\n",
    "            self.hidden = tuple([h.clone().detach() for h in self.hidden])\n",
    "        else:\n",
    "            self.hidden = self.hidden.clone().detach()\n",
    "\n",
    "    def reset_hidden(self, batch_size):\n",
    "        self.hidden = (torch.zeros(self.input_size, batch_size, self.hidden_size),\n",
    "                       torch.zeros(self.input_size, batch_size, self.hidden_size))\n",
    "\n",
    "    # train_epoch runs one epoch of training\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 = torch.randperm(input_data.shape[1])\n",
    "\n",
    "        self.reset_hidden(bs)\n",
    "\n",
    "        ep_loss = 0\n",
    "        #Iterate over batches of {bs} batches\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",
    "            # Load batch of shuffled segments\n",
    "            input_batch = input_data[:, shuffle[batch_i * bs:(batch_i + 1) * bs], :]\n",
    "            target_batch = target_data[:, shuffle[batch_i * bs:(batch_i + 1) * bs], :]\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",
    "            self.zero_grad()\n",
    "\n",
    "            # Choose the starting index for processing the rest of the batch sequence, in chunks of args.up_fr\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[0] - init_len) / up_fr)):\n",
    "                # Process input batch with neural network\n",
    "                output = self(input_batch[start_i:start_i + up_fr, :, :])\n",
    "\n",
    "                # Calculate loss and update network parameters\n",
    "                loss = loss_fcn(output, target_batch[start_i:start_i + up_fr, :, :])\n",
    "                loss.backward()\n",
    "                optim.step()\n",
    "\n",
    "                # Set the network hidden state, to detach it from the computation graph\n",
    "                # Die Gradienteninformation die einhergeht mit dem hidden state ist connected zu dem\n",
    "                # Computiaonal Graf von den vorherigen outputs. Da wir immer nur den letzten hidden state\n",
    "                # zur berechnung brauchen und den state auch nicht updaten wollen - wird er hier detached\n",
    "                # detached = wir löchen die Gradienteninformation\n",
    "                # https://discuss.pytorch.org/t/stupid-question-why-do-you-have-to-detach-the-hidden-state-of-lstms-but-not-the-hidden-state-of-a-linear-network/95089/3\n",
    "                self.detach_hidden()\n",
    "                self.zero_grad()\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": null,
   "id": "cbd60e7e671135d4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T13:48:32.467012700Z",
     "start_time": "2024-01-21T13:48:32.437598800Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class ESRLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ESRLoss, self).__init__()\n",
    "        self.epsilon = 0.00001\n",
    "\n",
    "    def forward(self, output, target):\n",
    "        loss = torch.add(target, -output)\n",
    "        loss = torch.pow(loss, 2)\n",
    "        loss = torch.mean(loss)\n",
    "        energy = torch.mean(torch.pow(target, 2)) + self.epsilon\n",
    "        loss = torch.div(loss, energy)\n",
    "        return loss\n",
    "class DCLoss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(DCLoss, self).__init__()\n",
    "        self.epsilon = 0.00001\n",
    "\n",
    "    def forward(self, output, target):\n",
    "        loss = torch.pow(torch.add(torch.mean(target, 0), -torch.mean(output, 0)), 2)\n",
    "        loss = torch.mean(loss)\n",
    "        energy = torch.mean(torch.pow(target, 2)) + self.epsilon\n",
    "        loss = torch.div(loss, energy)\n",
    "        return loss\n",
    "    \n",
    "class LossWrapper(nn.Module):\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(torch.Tensor([items[1] for items in loss_functions]))\n",
    "        except IndexError:\n",
    "            self.loss_factors = torch.ones(len(self.loss_functions))\n",
    "\n",
    "    def forward(self, output, target):\n",
    "        loss = 0\n",
    "        for i, losses in enumerate(self.loss_functions):\n",
    "            loss += torch.mul(losses(output, target), self.loss_factors[i])\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca470a4e37d92cd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T23:26:37.120936300Z",
     "start_time": "2024-01-21T23:26:37.101707400Z"
    },
    "collapsed": false
   },
   "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": null,
   "id": "ddc03fcf5a2209c3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-21T23:32:51.094930200Z",
     "start_time": "2024-01-21T23:26:39.885741800Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    print('CUDA device available')\n",
    "    torch.set_default_dtype(torch.cuda.FloatTensor)\n",
    "    torch.cuda.set_device(0)\n",
    "else:\n",
    "    print('CUDA device not available')\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",
    "\n",
    "optimiser = torch.optim.Adam(network.parameters(), lr=config[\"learning_rate\"], weight_decay=1e-4)\n",
    "# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimiser, 'min', factor=0.5, patience=5, verbose=True)\n",
    "# loss_functions = torch.nn.MSELoss(reduction='sum')\n",
    "loss_functions = LossWrapper(config[\"loss_fcns\"])\n",
    "\n",
    "summary = torchinfo.summary(network, (1, 1, 1), device=torch.device(\"cpu\"))\n",
    "print(summary)\n",
    "\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": "35e2061f",
   "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'].swapaxes(0, 1),\n",
    "                                     dataset.subsets['train']['target'].swapaxes(0, 1),\n",
    "                                     loss_functions,\n",
    "                                     optimiser,\n",
    "                                     config['batch_size'],\n",
    "                                     config['init_length'],\n",
    "                                     config['up_fr'])\n",
    "\n",
    "    # scheduler.step(epoch_loss)\n",
    "\n",
    "    print(\"Epoch loss:\", epoch_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0b75b4466489f46",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "name = \"model_0\"\n",
    "network.reset_hidden(1)\n",
    "torch.save(network.state_dict(), \"models/\"+name+\"/\"+\"stateful-lstm.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d058091229b579",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model = StatefulLSTM(input_size=config[\"input_size\"], \n",
    "                           output_size=config[\"output_size\"], \n",
    "                           hidden_size=config[\"hidden_size\"], \n",
    "                           skip=config[\"skip_con\"])\n",
    "\n",
    "model.load_state_dict(torch.load(\"models/\"+name+\"/\"+\"stateful-lstm.pth\", map_location=torch.device('cpu')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e71650307748c8f",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "input = dataset.subsets['train']['input'].swapaxes(0, 1)\n",
    "print(input.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "535228f882ad0aa4",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "output = torch.zeros_like(input)\n",
    "model.reset_hidden(1)\n",
    "with (torch.no_grad()):\n",
    "    for i in range(0, input.shape[1]):\n",
    "        output[:, i:i+1, :] = model(input[:, i:i+1, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87726a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_concat = output[:, 0, :]\n",
    "input_concat = input[:, 0, :]\n",
    "for i in range(1, output.shape[1]):\n",
    "    output_concat = torch.cat((output_concat, output[:, i, :]), 0)\n",
    "    input_concat = torch.cat((input_concat, input[:, i, :]), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2423babc043a6b17",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "write(os.path.join(\"input-pytorch.wav\"), 44100, input_concat.cpu().numpy().reshape(-1, 1))\n",
    "write(os.path.join(\"output-pytorch.wav\"), 44100, output_concat.cpu().numpy().reshape(-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5d300b88",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-08T14:42:14.055082700Z",
     "start_time": "2024-03-08T14:42:14.011276700Z"
    }
   },
   "outputs": [],
   "source": [
    "name = \"model_0\"\n",
    "model.reset_hidden(1)\n",
    "model_scripted = torch.jit.script(model)\n",
    "model_scripted.save(\"models/\"+name+\"/\"+\"stateful-lstm-dynamic.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "523dd9a8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-08T14:42:16.722936200Z",
     "start_time": "2024-03-08T14:42:16.683200100Z"
    }
   },
   "outputs": [],
   "source": [
    "example_input = torch.rand(2048, 1, 1)\n",
    "filepath = \"models/\"+name+\"/\"+\"stateful-lstm-libtorch-2048.onnx\"\n",
    "model.reset_hidden(1)\n",
    "\n",
    "# Export the model\n",
    "torch.onnx.export(model=model,\n",
    "                args=example_input,\n",
    "                f=filepath,\n",
    "                export_params=True,\n",
    "                opset_version=17,\n",
    "                do_constant_folding=True,\n",
    "                input_names = ['input'],\n",
    "                output_names = ['output'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd84631b080cd4d",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loaded_model = torch.jit.load(\"models/\"+name+\"/\"+\"stateful-lstm.pt\")\n",
    "\n",
    "# Example input\n",
    "sequence_length = 100  # adjust as per your model's training\n",
    "batch_size = 1         # can be set to 1 for testing individual sequences\n",
    "input_size = loaded_model.input_size  # should match the model's expected input size\\\n",
    "print(input_size)\n",
    "\n",
    "test_input = torch.zeros(2048, 1, 1)\n",
    "test_output = loaded_model(test_input)\n",
    "print(test_output)\n",
    "test_output = loaded_model(test_input)\n",
    "print(test_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad95a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_input = torch.zeros(2048, 1, 1)\n",
    "test_output = loaded_model(test_input)\n",
    "print(test_output)\n",
    "test_output = loaded_model(test_input)\n",
    "print(test_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4a5336bb61971ae",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "import numpy as np\n",
    "\n",
    "# Load the ONNX model\n",
    "onnx_model_path = \"models/\"+name+\"/\"+\"stateful-lstm-libtorch.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((2048, 1, 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)\n"
   ]
  }
 ],
 "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
}