SNGAN_Generators_32_64.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic knowledge\n",
    "\n",
    "Hereafter, I use `Padding.same` by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "let x1 = Tensor<Float>(randomNormal: [10, 32, 32, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let c3x3same = Conv2D<Float>(filterShape: (3, 3, 3, 32), strides: (1, 1), padding: Padding.same)\n",
    "let c3x3valid = Conv2D<Float>(filterShape: (3, 3, 3, 32), strides: (1, 1), padding: Padding.valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 32, 32, 32] [10, 30, 30, 32]\r\n"
     ]
    }
   ],
   "source": [
    "let y1 = c3x3same(x1)\n",
    "let y2 = c3x3valid(x1)\n",
    "\n",
    "print(y1.shape, y2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "let c1x1same = Conv2D<Float>(filterShape: (1, 1, 3, 32), strides: (1, 1), padding: Padding.same)\n",
    "let c1x1valid = Conv2D<Float>(filterShape: (1, 1, 3, 32), strides: (1, 1), padding: Padding.valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 32, 32, 32] [10, 32, 32, 32]\r\n"
     ]
    }
   ],
   "source": [
    "print(c1x1same(x1).shape, c1x1valid(x1).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "let unpool = UpSampling2D<Float>(size: 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 64, 64, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(unpool(x1).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "let avgval = AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2), padding: Padding.valid)\n",
    "let avgsame = AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2), padding: Padding.same)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 16, 16, 3] [10, 16, 16, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(avgval(x1).shape, avgsame(x1).shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "func conv<Scalar: TensorFlowFloatingPoint>(_ cIn: Int, _ cOut: Int, _ ks: Int = 3) -> Conv2D<Scalar> {\n",
    "    return Conv2D<Scalar>(filterShape: (ks, ks, cIn, cOut), strides: (1, 1), padding: Padding.same)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "func bn<Scalar: TensorFlowFloatingPoint>(_ nFeatures: Int) -> BatchNorm<Scalar> {\n",
    "    return BatchNorm<Scalar>(featureCount: nFeatures)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct ResBlock<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    \n",
    "    /// residual branch: bn1 -> act -> c1 -> bn2 -> act -> c2\n",
    "    public var bn1, bn2: BatchNorm<Scalar>\n",
    "    public var c1, c2: Conv2D<Scalar>\n",
    "\n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "\n",
    "    \n",
    "    init(_ cIn: Int, _ cOut: Int, _ cHidden: Int = -1, activation: @escaping Activation = relu) {\n",
    "        self.c1 = conv(cIn, cHidden)\n",
    "        self.c2 = conv(cHidden, cOut)\n",
    "        self.bn1 = bn(cIn)\n",
    "        self.bn2 = bn(cHidden)\n",
    "        self.activation = activation\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        let x = input\n",
    "        var h = bn1(input)\n",
    "        h = activation(h)\n",
    "        h = c1(h)\n",
    "        h = bn2(h)\n",
    "        h = activation(h)\n",
    "        h = c2(h)\n",
    "        return x + h\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 32, 32, 3] [10, 32, 32, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(ResBlock<Float>(3, 3, 3)(x1).shape, x1.shape)  // Assumed to be x1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct ResBlockWithConv<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    /// residual branch: bn1 -> act -> c1 -> bn2 -> act -> c2\n",
    "    public var bn1, bn2: BatchNorm<Scalar>\n",
    "    public var c1, c2: Conv2D<Scalar>\n",
    "    \n",
    "    /// body\n",
    "    public var cSC: Conv2D<Scalar>\n",
    "\n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "\n",
    "    \n",
    "    init(_ cIn: Int, _ cOut: Int, _ cHidden: Int = -1, activation: @escaping Activation = relu) {\n",
    "        self.c1 = conv(cIn, cHidden)\n",
    "        self.c2 = conv(cHidden, cOut)\n",
    "        self.cSC = conv(cIn, cOut, 1)\n",
    "        self.bn1 = bn(cIn)\n",
    "        self.bn2 = bn(cHidden)\n",
    "        self.activation = activation\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        let x = input\n",
    "        var h = bn1(input)\n",
    "        h = activation(h)\n",
    "        h = c1(h)\n",
    "        h = bn2(h)\n",
    "        h = activation(h)\n",
    "        h = c2(h)\n",
    "        return cSC(x) + h\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 32, 32, 32] [10, 32, 32, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(ResBlockWithConv<Float>(3, 32, 16)(x1).shape, x1.shape)  // Assumed to be [10, 32, 32, 32]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct ResBlockUpSample<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    /// residual branch: bn1 -> act -> upsample -> c1 -> bn2 -> act -> c2\n",
    "    public var bn1, bn2: BatchNorm<Scalar>\n",
    "    public var c1, c2: Conv2D<Scalar>\n",
    "    public var upsample: UpSampling2D<Scalar>\n",
    "\n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "\n",
    "    \n",
    "    init(_ cIn: Int, _ cOut: Int, _ cHidden: Int = -1, activation: @escaping Activation = relu) {\n",
    "        self.upsample = UpSampling2D<Scalar>(size: 2)\n",
    "        self.c1 = conv(cIn, cHidden)\n",
    "        self.c2 = conv(cHidden, cOut)\n",
    "        self.bn1 = bn(cIn)\n",
    "        self.bn2 = bn(cHidden)\n",
    "        self.activation = activation\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        let x = input\n",
    "        var h = bn1(input)\n",
    "        h = activation(h)\n",
    "        h = upsample(h)\n",
    "        h = c1(h)\n",
    "        h = bn2(h)\n",
    "        h = activation(h)\n",
    "        h = c2(h)\n",
    "        return upsample(x) + h\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 64, 64, 3] [10, 32, 32, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(ResBlockUpSample<Float>(3, 3, 3)(x1).shape, x1.shape)  // Assumed to be [10, 64, 64, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct ResBlockUpSampleWithConv<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    /// residual branch: bn1 -> act -> upsample -> c1 -> bn2 -> act -> c2\n",
    "    public var bn1, bn2: BatchNorm<Scalar>\n",
    "    public var c1, c2: Conv2D<Scalar>\n",
    "    public var upsample: UpSampling2D<Scalar>\n",
    "    \n",
    "    public var cSC: Conv2D<Scalar>\n",
    "\n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "\n",
    "    \n",
    "    init(_ cIn: Int, _ cOut: Int, _ cHidden: Int = -1, activation: @escaping Activation = relu) {\n",
    "        self.upsample = UpSampling2D<Scalar>(size: 2)\n",
    "        self.c1 = conv(cIn, cHidden)\n",
    "        self.c2 = conv(cHidden, cOut)\n",
    "        self.cSC = conv(cIn, cOut, 1)\n",
    "        self.bn1 = bn(cIn)\n",
    "        self.bn2 = bn(cHidden)\n",
    "        self.activation = activation\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        let x = input\n",
    "        var h = bn1(input)\n",
    "        h = activation(h)\n",
    "        h = upsample(h)\n",
    "        h = c1(h)\n",
    "        h = bn2(h)\n",
    "        h = activation(h)\n",
    "        h = c2(h)\n",
    "        return cSC(upsample(x)) + h\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[10, 64, 64, 32] [10, 32, 32, 3]\r\n"
     ]
    }
   ],
   "source": [
    "print(ResBlockUpSampleWithConv<Float>(3, 32, 32)(x1).shape, x1.shape)  // Assumed to be [10, 64, 64, 32]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Generator32<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    public var l1: Dense<Scalar>\n",
    "    public var block2, block3, block4: ResBlockUpSample<Scalar>\n",
    "    public var bn5: BatchNorm<Scalar>\n",
    "    public var c5: Conv2D<Scalar>\n",
    "    \n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "    \n",
    "    @noDerivative public let nFeatures: Int\n",
    "    @noDerivative public let bottomWidth: Int\n",
    "    \n",
    "    init(_ nFeatures: Int = 256, _ dimZ: Int = 128, _ bottomWidth: Int = 4, activation: @escaping Activation = relu) {\n",
    "        self.l1 = Dense<Scalar>(inputSize: dimZ, outputSize: bottomWidth * bottomWidth * nFeatures)\n",
    "        self.block2 = ResBlockUpSample<Scalar>(nFeatures, nFeatures, nFeatures)\n",
    "        self.block3 = ResBlockUpSample<Scalar>(nFeatures, nFeatures, nFeatures)\n",
    "        self.block4 = ResBlockUpSample<Scalar>(nFeatures, nFeatures, nFeatures)\n",
    "        self.bn5 = bn(nFeatures)\n",
    "        self.c5 = conv(nFeatures, 3)\n",
    "        \n",
    "        self.activation = activation\n",
    "        \n",
    "        self.nFeatures = nFeatures\n",
    "        self.bottomWidth = bottomWidth\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        var h = l1(input)\n",
    "        h = h.reshaped(to: TensorShape([h.shape[0], bottomWidth, bottomWidth, nFeatures]))\n",
    "        h = block2(h)\n",
    "        h = block3(h)\n",
    "        h = block4(h)\n",
    "        h = bn5(h)\n",
    "        h = activation(h)\n",
    "        h = c5(h)\n",
    "        return tanh(h)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ [10, 32, 32, 3]\n",
       "  ▿ dimensions : 4 elements\n",
       "    - 0 : 10\n",
       "    - 1 : 32\n",
       "    - 2 : 32\n",
       "    - 3 : 3\n"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Generator32<Float>()(Tensor<Float>(randomUniform: [10, 128])).shape  // Assumed to be (10, 32, 32, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Generator64<Scalar: TensorFlowFloatingPoint>: Layer {\n",
    "    public var l1: Dense<Scalar>\n",
    "    public var block2, block3, block4, block5: ResBlockUpSampleWithConv<Scalar>\n",
    "    public var bn6: BatchNorm<Scalar>\n",
    "    public var c6: Conv2D<Scalar>\n",
    "    \n",
    "    @noDerivative public let activation: Activation\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "    \n",
    "    @noDerivative public let nFeatures: Int\n",
    "    @noDerivative public let bottomWidth: Int\n",
    "    \n",
    "    init(_ nFeatures: Int = 64, _ dimZ: Int = 128, _ bottomWidth: Int = 4, activation: @escaping Activation = relu) {\n",
    "        self.l1 = Dense<Scalar>(inputSize: dimZ, outputSize: bottomWidth * bottomWidth * nFeatures * 16)\n",
    "        self.block2 = ResBlockUpSampleWithConv<Scalar>(nFeatures * 16, nFeatures * 8, nFeatures * 8)\n",
    "        self.block3 = ResBlockUpSampleWithConv<Scalar>(nFeatures * 8, nFeatures * 4, nFeatures * 4)\n",
    "        self.block4 = ResBlockUpSampleWithConv<Scalar>(nFeatures * 4, nFeatures * 2, nFeatures * 2)\n",
    "        self.block5 = ResBlockUpSampleWithConv<Scalar>(nFeatures * 2, nFeatures, nFeatures)\n",
    "        self.bn6 = bn(nFeatures)\n",
    "        self.c6 = conv(nFeatures, 3)\n",
    "        \n",
    "        self.activation = activation\n",
    "        \n",
    "        self.nFeatures = nFeatures\n",
    "        self.bottomWidth = bottomWidth\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        var h = l1(input)\n",
    "        h = h.reshaped(to: TensorShape([h.shape[0], bottomWidth, bottomWidth, nFeatures * 16]))\n",
    "        h = block2(h)\n",
    "        h = block3(h)\n",
    "        h = block4(h)\n",
    "        h = block5(h)\n",
    "        h = bn6(h)\n",
    "        h = activation(h)\n",
    "        h = c6(h)\n",
    "        return tanh(h)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "▿ [10, 64, 64, 3]\n",
       "  ▿ dimensions : 4 elements\n",
       "    - 0 : 10\n",
       "    - 1 : 64\n",
       "    - 2 : 64\n",
       "    - 3 : 3\n"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Generator64<Float>()(Tensor<Float>(randomUniform: [10, 128])).shape  // Assumed to be (10, 64, 64, 3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Swift",
   "language": "swift",
   "name": "swift"
  },
  "language_info": {
   "file_extension": ".swift",
   "mimetype": "text/x-swift",
   "name": "swift",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}