Skip to content

Instantly share code, notes, and snippets.

@crcrpar
Created July 26, 2019 17:35
Show Gist options
  • Select an option

  • Save crcrpar/964a0f3512a0aa2944851c7f5b3f797f to your computer and use it in GitHub Desktop.

Select an option

Save crcrpar/964a0f3512a0aa2944851c7f5b3f797f to your computer and use it in GitHub Desktop.
Swift for TensorFlow implementation of SNConv2D, SNDense, Blocks, and Models
Display the source blob
Display the rendered blob
Raw
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Display the source blob
Display the rendered blob
Raw
{
"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