facebook_model_split_integrado.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import LSTM\n",
    "import torchaudio\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#network parameters set @ facebook\n",
    "L = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 64000])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#real audio example\n",
    "audio_input, s_rate = torchaudio.load('/home/minds/Desktop/test_speech_mix/generated/speaker/id10009/7hpSiT9_gCE/0.wav', normalization=True)\n",
    "audio_input = audio_input.unsqueeze(0)\n",
    "audio_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseRNNBlock(nn.Module):\n",
    "    '''\n",
    "    O bloco impar passa a RNN ao longo da dimensão do tempo, ou seja,\n",
    "    ao longo de R chunks de tamanho K\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, Parameter=128, hidden_size=128, **kwargs):\n",
    "        super(BaseRNNBlock, self).__init__()\n",
    "        self.lstm1 = LSTM(input_size=Parameter, hidden_size=hidden_size, batch_first=True, bidirectional=True, **kwargs)\n",
    "        self.lstm2 = LSTM(input_size=Parameter, hidden_size=hidden_size, batch_first=True, bidirectional=True, **kwargs)\n",
    "        #no paper P\n",
    "        self.P = nn.Linear(hidden_size*2 + Parameter, Parameter)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        outp1, _ = self.lstm1(x)\n",
    "        outp2, _ = self.lstm2(x)\n",
    "        outp = torch.cat((outp1 * outp2, x), dim=-1)\n",
    "        #O tensor volta a NxKxR\n",
    "        return self.P(outp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNBlock(nn.Module):\n",
    "    '''\n",
    "    Contem um bloco par e um bloco impar\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, K, R, hidden_size=128,**kwargs):\n",
    "        super(RNNBlock, self).__init__()\n",
    "        self.oddBlock = BaseRNNBlock(Parameter=K, hidden_size=hidden_size, **kwargs)\n",
    "        self.evenBlock = BaseRNNBlock(Parameter=R, hidden_size=hidden_size, **kwargs)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        outp = self.oddBlock(x)\n",
    "        #skip connection\n",
    "        x += outp\n",
    "        #Tensor NxRxK -> NxKxR\n",
    "        x = torch.transpose(x, 1,-1)\n",
    "        outp = self.evenBlock(x)\n",
    "        #skip connection\n",
    "        x += outp\n",
    "        #Tensor NxKxR -> NxRxK\n",
    "        x = torch.transpose(x, 1,-1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Facebookmodel(nn.Module):\n",
    "    '''\n",
    "    Modelo do facebook\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, n, k, r, c=2, l=8, b=1, **kwargs):\n",
    "        super(Facebookmodel, self).__init__()\n",
    "        assert l % 2 == 0, 'l must be even'\n",
    "        self.c = c\n",
    "        self.encoder = nn.Conv1d(1, n, l, int(l/2))\n",
    "        self.rnnblocks = [RNNBlock(k, r, **kwargs) for _ in range(b)]\n",
    "        self.d = nn.Conv1d(r, c*r, kernel_size=1)\n",
    "        self.activation = torch.nn.PReLU(num_parameters=1, init=0.25)\n",
    "        self.decoder = nn.ConvTranspose1d(n, 1, kernel_size=l, stride=int(l/2))\n",
    "        \n",
    "        #teste de decode com uma convolucao 2d\n",
    "        self.decoder2d = nn.ConvTranspose2d(n, 1, kernel_size=(1,l), stride=(1,int(l/2)))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        encoded = self.encoder(x).squeeze(0)\n",
    "        chunks = self.chunk(encoded)\n",
    "        print(chunks.shape)\n",
    "        outps = list()\n",
    "        for block in self.rnnblocks:\n",
    "            chunks = block(chunks)\n",
    "            res = self.d(self.activation(chunks))\n",
    "            outps.append(res)\n",
    "        \n",
    "        outps = self.apply_overlap_and_add(outps)\n",
    "        return self.decode2d(outps)\n",
    "        #return self.decode(o1, o2)\n",
    "    \n",
    "    def chunk(self, x):\n",
    "        x = torch.cat((x, torch.zeros((64, 110))), dim=-1)\n",
    "        x = torch.cat((torch.zeros((64, 89)), x), dim=-1)\n",
    "        \n",
    "        return x.unfold(-1, 178, 89)\n",
    "        \n",
    "    def decode2d(self, x):\n",
    "        '''\n",
    "        Testar o decode com uma convolucao de 2 dimensoes, audios das c fontes\n",
    "        juntos no mesmo tensor\n",
    "        '''\n",
    "        \n",
    "        restored = []\n",
    "        for a in x:\n",
    "            a = a[...,89:-110].unsqueeze(0)\n",
    "            d = self.decoder2d(a).squeeze(1)\n",
    "            restored.append(d)\n",
    "        \n",
    "        return restored\n",
    "    \n",
    "    def decode(self, x):\n",
    "        '''\n",
    "        Decode de separacao com convolucao 1d, os audios das c fontes diferentes sao separados\n",
    "        previamente\n",
    "        '''\n",
    "        restored = [[] for _ in range(self.c)]\n",
    "        \n",
    "        for a in x:\n",
    "            for i in range(self.c):\n",
    "                t = a[:,i,89:-110].unsqueeze(0)\n",
    "                t = self.decoder(t)\n",
    "                restored[i].append(t.squeeze(0))\n",
    "        \n",
    "        return restored\n",
    "    \n",
    "    def apply_overlap_and_add(self, x):\n",
    "        overlapped_added = list()\n",
    "        for el in x:\n",
    "            result = self.overlap_and_add(el)\n",
    "            overlapped_added.append(result)\n",
    "            \n",
    "        return overlapped_added\n",
    "    \n",
    "    def overlap_and_add(self, x):\n",
    "        '''\n",
    "        Faz overlap and add usando pytorch fold\n",
    "        '''\n",
    "        x = torch.transpose(x, -2, -1)\n",
    "        result = torch.nn.functional.fold(x, (self.c, 16198) ,kernel_size=(1,178), stride=(1,89))\n",
    "        return result.squeeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = Facebookmodel(64, 178, 181, b = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 181, 178])\n"
     ]
    }
   ],
   "source": [
    "r = m(audio_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 64000])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.rand((1,2,64000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [],
   "source": [
    " t = torch.rand((1,64, 3, 15999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = nn.ConvTranspose2d(64, 1, kernel_size=(1,8), stride=(1,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 3, 64000])"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m(t)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = e(audio_input).squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "dec = nn.ConvTranspose1d(64, 1, kernel_size=L, stride=int(L/2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 64000])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dec(t.unsqueeze(0)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 15999])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.2963,  0.2988,  0.2976,  ...,  0.3097,  0.2949,  0.2941],\n",
       "        [ 0.0497,  0.0493,  0.0543,  ...,  0.0472,  0.0316,  0.0453],\n",
       "        [-0.2260, -0.2260, -0.2255,  ..., -0.2244, -0.2102, -0.2111],\n",
       "        ...,\n",
       "        [-0.3466, -0.3421, -0.3494,  ..., -0.3353, -0.3397, -0.3497],\n",
       "        [-0.1494, -0.1512, -0.1506,  ..., -0.1456, -0.1427, -0.1534],\n",
       "        [-0.2383, -0.2415, -0.2306,  ..., -0.2481, -0.2472, -0.2326]],\n",
       "       grad_fn=<SqueezeBackward1>)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 16198])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t0 = torch.cat((t, torch.zeros((64, 110))), dim=-1)\n",
    "t0 = torch.cat((torch.zeros((64, 89)), t0), -1)\n",
    "t0.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.2963,  0.2988,  0.2976,  ...,  0.3097,  0.2949,  0.2941],\n",
       "        [ 0.0497,  0.0493,  0.0543,  ...,  0.0472,  0.0316,  0.0453],\n",
       "        [-0.2260, -0.2260, -0.2255,  ..., -0.2244, -0.2102, -0.2111],\n",
       "        ...,\n",
       "        [-0.3466, -0.3421, -0.3494,  ..., -0.3353, -0.3397, -0.3497],\n",
       "        [-0.1494, -0.1512, -0.1506,  ..., -0.1456, -0.1427, -0.1534],\n",
       "        [-0.2383, -0.2415, -0.2306,  ..., -0.2481, -0.2472, -0.2326]],\n",
       "       grad_fn=<SliceBackward>)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t0[:,89:-110]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 181, 178])"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = t0.unfold(-1,178,89)\n",
    "t.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "fold = nn.Fold(output_size=(1, 16198), kernel_size=(1, 178), stride=(1, 178))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Given output_size=(1, 16198), kernel_size=(1, 178), dilation=(1, 1), padding=(0, 0), stride=(1, 178), expected size of input's dimension 2 to match the calculated number of sliding blocks 1 * 91 = 91, but got input.size(2)=181.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-60-67a69db52b7b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfold\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    548\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    549\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 550\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    551\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    552\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/modules/fold.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    128\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    129\u001b[0m         return F.fold(input, self.output_size, self.kernel_size, self.dilation,\n\u001b[0;32m--> 130\u001b[0;31m                       self.padding, self.stride)\n\u001b[0m\u001b[1;32m    131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    132\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mfold\u001b[0;34m(input, output_size, kernel_size, dilation, padding, stride)\u001b[0m\n\u001b[1;32m   3660\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3661\u001b[0m         return torch._C._nn.col2im(input, _pair(output_size), _pair(kernel_size),\n\u001b[0;32m-> 3662\u001b[0;31m                                    _pair(dilation), _pair(padding), _pair(stride))\n\u001b[0m\u001b[1;32m   3663\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3664\u001b[0m         \u001b[0;32mraise\u001b[0m \u001b[0mNotImplementedError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input Error: Only 3D input Tensors are supported (got {}D)\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Given output_size=(1, 16198), kernel_size=(1, 178), dilation=(1, 1), padding=(0, 0), stride=(1, 178), expected size of input's dimension 2 to match the calculated number of sliding blocks 1 * 91 = 91, but got input.size(2)=181."
     ]
    }
   ],
   "source": [
    "res = fold(torch.transpose(t, -1,-2)).squeeze(1).squeeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'res' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-61-9341671a9fa9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;34m(\u001b[0m\u001b[0mt0\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'res' is not defined"
     ]
    }
   ],
   "source": [
    "(t0 == res).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import functional as nnf\n",
    "result = torch.tensor([[1,2,4,6.,7], [1,2,4,6,7]]).unsqueeze(0).unsqueeze(0)\n",
    "recovered = nnf.unfold(result, kernel_size=(2,3), stride=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = torch.rand((1,4,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = nn.Fold(output_size=(1,4), kernel_size=(1,4), stride=(1,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f(tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.6204, 0.2541, 0.5079,  ..., 0.5675, 0.3362, 0.8327]])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = torch.rand((1,16198))\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[0.6204, 0.2541, 0.5079,  ..., 0.4269, 0.8965, 0.3063],\n",
       "          [0.8041, 0.0994, 0.6862,  ..., 0.4787, 0.2633, 0.9002],\n",
       "          [0.5556, 0.5323, 0.0520,  ..., 0.6909, 0.8630, 0.1000],\n",
       "          ...,\n",
       "          [0.5980, 0.1780, 0.8383,  ..., 0.0549, 0.7329, 0.6425],\n",
       "          [0.2111, 0.6527, 0.5950,  ..., 0.7827, 0.4977, 0.7434],\n",
       "          [0.2652, 0.1109, 0.0736,  ..., 0.5675, 0.3362, 0.8327]]]),\n",
       " torch.Size([1, 181, 178]))"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f = t.unfold(-1, 178, 89)\n",
    "f, f.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 178, 181])"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.transpose(f, -2, -1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.6204, 0.2541, 0.5079,  ..., 0.5675, 0.3362, 0.8327]]]])"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = nnf.fold(torch.transpose(f, -2, -1), (1, 16198), kernel_size=(1, 178), stride=(1, 89))\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(res.squeeze(1).squeeze(1) == t).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recovered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.tensor([[1, 2, 3], [1, 2, 3.]])    \n",
    "b = torch.tensor([[5, 6, 7], [5, 6, 7.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uf = torch.cat((a.view(1, 6, 1), b.view(1, 6, 1)), dim=2)\n",
    "uf, uf.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = nnf.fold(uf, (2,5), kernel_size=(2,3), stride=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}