Fix code formatting

Author: vvmnnnkv

FastStyleTransfer/Demo/Helpers.swift

@@ -2,10 +2,10 @@ import Foundation
 import TensorFlow
 import FastStyleTransfer
 
-// Make model importable
 extension TransformerNet: ImportableLayer {}
 
-func parseArgs<T>(into obj: inout T, with params: [String: WritableKeyPath<T, String?>]) {
+/// Updates `obj` with values from command line arguments according to `params` map.
+func parseArguments<T>(into obj: inout T, with params: [String: WritableKeyPath<T, String?>]) {
     for arg in CommandLine.arguments.dropFirst() {
         if !arg.starts(with: "--") { continue }
         let parts = arg.split(separator: "=", maxSplits: 2)
@@ -17,15 +17,15 @@ func parseArgs<T>(into obj: inout T, with params: [String: WritableKeyPath<T, St
 }
 
 enum FileError: Error {
-    case file_not_found
+    case fileNotFound
 }
 
+/// Updates `model` with parameters from numpy archive file in `path`.
 func importWeights(_ model: inout TransformerNet, from path: String) throws {
     guard FileManager.default.fileExists(atPath: path) else {
-        throw FileError.file_not_found
+        throw FileError.fileNotFound
     }
-    // Map of model params to loaded params
-    // Names don't match exactly, and axes in filters need to be reversed
+    // Names don't match exactly, and axes in filters need to be reversed.
     let map = [
         "conv1.conv2d.filter": ("conv1.conv2d.weight", [3, 2, 1, 0]),
         "conv2.conv2d.filter": ("conv2.conv2d.weight", [3, 2, 1, 0]),
@@ -77,14 +77,16 @@ func importWeights(_ model: inout TransformerNet, from path: String) throws {
     model.unsafeImport(fromNumpyArchive: path, map: map)
 }
 
+/// Loads from `file` and returns JPEG image as HxWxC tensor of floats in (0..1) range.
 func loadJpegAsTensor(from file: String) throws -> Tensor<Float> {
     guard FileManager.default.fileExists(atPath: file) else {
-        throw FileError.file_not_found
+        throw FileError.fileNotFound
     }
     let imgData = Raw.readFile(filename: StringTensor(file))
     return Tensor<Float>(Raw.decodeJpeg(contents: imgData, channels: 3, dctMethod: "")) / 255
 }
 
+/// Clips & converts HxWxC `tensor` of floats to byte range and saves as JPEG.
 func saveTensorAsJpeg(_ tensor: Tensor<Float>, to file: String) {
     let clipped = Raw.clipByValue(t: tensor, clipValueMin: Tensor(0), clipValueMax: Tensor(255))
     let jpg = Raw.encodeJpeg(image: Tensor<UInt8>(clipped), format: .rgb, xmpMetadata: "")

FastStyleTransfer/Demo/main.swift

@@ -11,31 +11,33 @@ func printUsage() {
     print("    --output: Path to output image")
 }
 
+/// Startup parameters.
 struct Config {
     var weights: String? = "Demo/weights/candy.npz"
     var image: String? = nil
     var output: String? = "out.jpg"
 }
+
 var config = Config()
-parseArgs(into: &config, with: [
+parseArguments(into: &config, with: [
     "weights": \Config.weights, 
     "image": \Config.image, 
-    "output": \Config.output])
+    "output": \Config.output
+])
 
 guard let image = config.image, let output = config.output else {
     print("Error: No input image!")
     printUsage()
     exit(1)
 }
 
-// load image
 guard let imageTensor = try? loadJpegAsTensor(from: image) else {
     print("Error: Failed to load image \(image). Check file exists and has JPEG format")
     printUsage()
     exit(1)
 }
 
-// init model
+// Init the model.
 var style = TransformerNet()
 do {
     try importWeights(&style, from: config.weights!)
@@ -45,6 +47,8 @@ do {
     exit(1)
 }
 
+// Apply the model to image.
 let out = style(imageTensor.expandingShape(at: 0))
+
 saveTensorAsJpeg(out.squeezingShape(at: 0), to: output)
-print("Written output to \(output)")
+print("Written output to \(output)")

FastStyleTransfer/Layers/Helpers.swift

@@ -1,20 +1,33 @@
 import TensorFlow
 
-/// Layer for padding with reflection over mini-batch of images
-/// Expected input layout is BxHxWxC
-public struct ReflectionPad2d<Scalar: TensorFlowFloatingPoint>: Layer {
+/// A 2-D layer applying padding with reflection over a mini-batch.
+public struct ReflectionPad2D<Scalar: TensorFlowFloatingPoint>: Layer {
+    /// The padding values along the spatial dimensions.
     @noDerivative public let padding: ((Int, Int), (Int, Int))
 
+    /// Creates a reflect-padding 2D Layer.
+    ///
+    /// - Parameter padding: A tuple of 2 tuples of two integers describing how many elements to
+    ///   be padded at the beginning and end of each padding dimensions.
     public init(padding: ((Int, Int), (Int, Int))) {
         self.padding = padding
     }
 
+    /// Creates a reflect-padding 2D Layer.
+    ///
+    /// - Parameter padding: Integer that describes how many elements to be padded
+    ///   at the beginning and end of each padding dimensions.
     public init(padding: Int) {
         self.padding = ((padding, padding), (padding, padding))
     }
 
+    /// Returns the output obtained from applying the layer to the given input.
+    ///
+    /// - Parameter input: The input to the layer. Expected layout is BxHxWxC.
+    /// - Returns: The output.
     @differentiable
     public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        // Padding applied to height and width dimensions only.
         return input.paddedWithReflection(forSizes: [
             (0, 0),
             padding.0,
@@ -24,11 +37,14 @@ public struct ReflectionPad2d<Scalar: TensorFlowFloatingPoint>: Layer {
     }
 }
 
-
-/// Layer applying relu activation function
+/// A layer applying `relu` activation function.
 public struct ReLU<Scalar: TensorFlowFloatingPoint>: Layer {
+    /// Returns the output obtained from applying the layer to the given input.
+    ///
+    /// - Parameter input: The input to the layer.
+    /// - Returns: The output.
     @differentiable
     public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
         return relu(input)
     }
-}
+}

FastStyleTransfer/Layers/Normalization.swift

@@ -1,21 +1,34 @@
 import TensorFlow
 
-/// Layer that applies instance normalization over a mini-batch of images
-/// Expected input layout is BxHxWxC
+/// 2-D layer applying instance normalization over a mini-batch of inputs.
+///
 /// Reference: [Instance Normalization](https://arxiv.org/abs/1607.08022)
-public struct InstanceNorm2d<Scalar: TensorFlowFloatingPoint>: Layer {
+public struct InstanceNorm2D<Scalar: TensorFlowFloatingPoint>: Layer {
+    /// Learnable parameter scale for affine transformation.
     public var scale: Tensor<Scalar>
+    /// Learnable parameter offset for affine transformation.
     public var offset: Tensor<Scalar>
+    /// Small value added in denominator for numerical stability.
     @noDerivative public var epsilon: Tensor<Scalar>
 
+    /// Creates a instance normalization 2D Layer.
+    ///
+    /// - Parameters:
+    ///   - featureCount: Size of the channel axis in the expected input.
+    ///   - epsilon: Small scalar added for numerical stability.
     public init(featureCount: Int, epsilon: Tensor<Scalar> = Tensor(1e-5)) {
         self.epsilon = epsilon
         scale = Tensor<Scalar>(ones: [featureCount])
         offset = Tensor<Scalar>(zeros: [featureCount])
     }
 
+    /// Returns the output obtained from applying the layer to the given input.
+    ///
+    /// - Parameter input: The input to the layer. Expected input layout is BxHxWxC.
+    /// - Returns: The output.
     @differentiable
     public func callAsFunction(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
+        // Calculate mean & variance along H,W axes.
         let mean = input.mean(alongAxes: [1, 2])
         let variance = input.variance(alongAxes: [1, 2])
         let norm = (input - mean) * rsqrt(variance + epsilon)