Created
July 26, 2019 17:35
-
-
Save crcrpar/964a0f3512a0aa2944851c7f5b3f797f to your computer and use it in GitHub Desktop.
Swift for TensorFlow implementation of SNConv2D, SNDense, Blocks, and Models
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| "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 | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment