Minor refactor and doc changes (#68)

Bender.xcodeproj/project.pbxproj

@@ -9,6 +9,7 @@
 /* Begin PBXBuildFile section */
 		28F828811C494B2C00330CF4 /* Bender.h in Headers */ = {isa = PBXBuildFile; fileRef = 28F828801C494B2C00330CF4 /* Bender.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		46354BC41EEF320700B083EF /* DependencyListBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = 46354BC31EEF320700B083EF /* DependencyListBuilder.swift */; };
+		4675C5561F4209DF00895175 /* Device.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4675C5551F4209DF00895175 /* Device.swift */; };
 		4697C2191ED31F8D003ED2E9 /* local_response_norm.metal in Sources */ = {isa = PBXBuildFile; fileRef = 4697C2181ED31F8D003ED2E9 /* local_response_norm.metal */; };
 		4697C21B1ED35E50003ED2E9 /* LocalResponseNorm.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4697C21A1ED35E50003ED2E9 /* LocalResponseNorm.swift */; };
 		46FAFCB91F029B09008754E1 /* Concat.swift in Sources */ = {isa = PBXBuildFile; fileRef = 46FAFCB81F029B09008754E1 /* Concat.swift */; };
@@ -91,6 +92,7 @@
 		28F828801C494B2C00330CF4 /* Bender.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = Bender.h; sourceTree = "<group>"; };
 		28F828821C494B2C00330CF4 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
 		46354BC31EEF320700B083EF /* DependencyListBuilder.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = DependencyListBuilder.swift; sourceTree = "<group>"; };
+		4675C5551F4209DF00895175 /* Device.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Device.swift; sourceTree = "<group>"; };
 		4697C2181ED31F8D003ED2E9 /* local_response_norm.metal */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.metal; path = local_response_norm.metal; sourceTree = "<group>"; };
 		4697C21A1ED35E50003ED2E9 /* LocalResponseNorm.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LocalResponseNorm.swift; sourceTree = "<group>"; };
 		46FAFCB81F029B09008754E1 /* Concat.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Concat.swift; sourceTree = "<group>"; };
@@ -269,6 +271,7 @@
 				8FA984001EBA70F900586D1B /* NetworkLayer.swift */,
 				8FA8E32F1EC0FE2700E8BAD8 /* Operators.swift */,
 				8F60A31C1EC4E6E400296167 /* PaddingType.swift */,
+				4675C5551F4209DF00895175 /* Device.swift */,
 			);
 			path = Core;
 			sourceTree = "<group>";
@@ -484,6 +487,7 @@
 				8FAB77221ECC86170050AB16 /* attr_value.pb.swift in Sources */,
 				8FFF4AD71EDEFDB600141E9D /* Shape.swift in Sources */,
 				8FA8E3301EC0FE2700E8BAD8 /* Operators.swift in Sources */,
+				4675C5561F4209DF00895175 /* Device.swift in Sources */,
 				8F79B8031EC23563002D02C8 /* ResidualLayer.swift in Sources */,
 				8FAB77291ECC86170050AB16 /* versions.pb.swift in Sources */,
 				8FA8E3211EC0B8B500E8BAD8 /* BGRAtoRGBA.swift in Sources */,

Documentation/API.md

@@ -16,8 +16,7 @@ This document explains the basic API in __Bender__.
 To create a network model you can create it from scratch or [import](Importing.md) it from a TensorFlow graph. We will explain how to create a network from scratch:
 
 ```swift
-let network = Network(device: device,
-                      inputSize: inputSize,
+let network = Network(inputSize: inputSize,
                       parameterLoader: loader)
 
 network.start
@@ -33,7 +32,7 @@ network.start
 network.initialize()
 ```
 
-First, we have to create the `network` which receives the MTLDevice (GPU), an inputSize and a parameter loader. The network comes with a `start` node which is the starting point of the network. The `inputSize` is the size expected by the first layer in the network. If the images you pass the network to be processed are not of the expected size then the `start` node will resize them accordingly.
+First, we have to create the `network` which receives the inputSize and a parameter loader. The network comes with a `start` node which is the starting point of the network. The `inputSize` is the size expected by the first layer in the network. If the images you pass the network to be processed are not of the expected size then the `start` node will resize them accordingly.
 
 The `parameterLoader` is responsible for loading the weights for each layer. It will be explained in detail further below.
 
@@ -54,16 +53,14 @@ After you finish adding layers to your network, you must call `network.initializ
 
 ## Running a network
 
-To run a network call `run(...)`:
+To run a network call `run(/* ... */)`:
 
 ```swift
-let commandQueue: MTLCommandQueue = ...
-
 // get image from somewhere
-let image = MPSImage(...)
+let image = MPSImage(/* ... */)
 
-network.run(inputImage: image, queue: commandQueue) { outputImage in
-    ...
+network.run(input: image) { output in
+    // ...
 }
 ```
 

Documentation/Importing.md

@@ -28,20 +28,16 @@ benderthon tf-freeze checkpoint_path.ckpt graph_with_weights.pb Output_Node_Name
 To import a model saved in a [Protobuf](https://developers.google.com/protocol-buffers/) file you must add it to your Xcode project and load it like this:
 
 ```swift
-// Define network
-network = Network(device: device, inputSize: inputSize, parameterLoader: nil)
 
-// Load graph file
+// Set an url pointing to your model
 let url = Bundle.main.url(forResource: "myGraph", withExtension: "pb")!
 
-// Create converter
+// Create the converter
 let converter = TFConverter.default()
 
-// Convert graph
-network.convert(converter: converter, url: url, type: .binary)
+// Load it
+let network = Network.load(url: url, converter: converter, inputSize: LayerSize(h: 256, w: 256, f: 3))
 
-// Initialize network
-network.initialize()
 ```
 
 `TFConverter` is the class responsible for converting a TF model to Bender. It will try to map nodes or groups of nodes in the TF graph to Bender layers. If it encounters unknown nodes then it will ignore them. This means that a graph might be disconnected if your TF model uses functions that are not implemented in Bender.
@@ -79,8 +75,7 @@ func optimize(graph: TFGraph)
 After you create the optimizer, you have to add it to your `TFConverter` like this:
 
 ```swift
-let converter = TFConverter.default()
-converter.optimizers.append(MyTFOptimizer())
+let converter = TFConverter.default(additionalOptimizers: [MyTFOptimizer()])
 ```
 
 #### Removing nodes

Example/Example/GrayScale.swift

@@ -27,7 +27,7 @@ class GrayScale: NetworkLayer {
         assert(getIncoming().count == 1, "GrayScale must have one input")
         let incoming = getIncoming()[0]
         assert(incoming.outputSize.f <= 4, "GrayScale input must have at most 4 feature channels")
-        outputSize = LayerSize(f: outputChannels, w: incoming.outputSize.w, h: incoming.outputSize.h)
+        outputSize = LayerSize(h: incoming.outputSize.h, w: incoming.outputSize.w, f: outputChannels)
         outputImage = MPSImage(device: device, imageDescriptor: MPSImageDescriptor(layerSize: outputSize))
     }
 

Example/Example/StyleTransferViewController.swift

@@ -15,17 +15,15 @@ class StyleTransferViewController: UIViewController, ExampleViewController {
 
     var styleNet: Network!
     var styleNet2: Network!
-    var device: MTLDevice!
     var commandQueue: MTLCommandQueue!
-    let inputSize = LayerSize(f: 3, w: 256)
+    let inputSize = LayerSize(h: 256, w: 256, f: 3)
 
     var pixelBufferPool: CVPixelBufferPool?
     @IBOutlet weak var imageView: UIImageView!
 
     override func viewDidLoad() {
         super.viewDidLoad()
-        self.device = MTLCreateSystemDefaultDevice()!
-        self.commandQueue = device.makeCommandQueue()
+        self.commandQueue = Device.shared.makeCommandQueue()
         setupStyleNet()
         var me = self
         me.setPixelBufferPool()
@@ -39,25 +37,24 @@ class StyleTransferViewController: UIViewController, ExampleViewController {
     func setupStyleNet() {
         measure("Set up takes:") {
 
-            styleNet = Network(device: device, inputSize: inputSize, parameterLoader: nil)
-
             let url = Bundle.main.url(forResource: "g_and_w2", withExtension: "pb")!
-            let converter = TFConverter.default()
-            converter.optimizers.append(TFInstanceNormOptimizer())
+            let converter = TFConverter.default(additionalOptimizers: [TFInstanceNormOptimizer()])
 
-            styleNet.convert(converter: converter, url: url, type: .binary)
+            styleNet = Network.load(url: url, converter: converter, inputSize: inputSize, performInitialize: false)
             styleNet.addPostProcessing(layers: [ImageLinearTransform()])
 
+            // after adding all our layers we are able to initialize the network
+
             styleNet.initialize()
         }
     }
 
     @IBAction func runNetwork(_ sender: Any) {
         let buffer = commandQueue.makeCommandBuffer()
-        let image = loadTestImage(device: device, commandBuffer: buffer)
+        let image = loadTestImage(commandBuffer: buffer)
         buffer.commit()
         buffer.waitUntilCompleted()
-        styleNet.run(inputImage: image, queue: commandQueue) { [weak self] imageA in
+        styleNet.run(input: image, queue: commandQueue) { [weak self] imageA in
             if let buffer = self?.getPixelBuffer(from: imageA.texture, bufferPool: self!.pixelBufferPool!) {
                 let ciImage = CIImage(cvImageBuffer: buffer)
                 let context = CIContext()
@@ -70,9 +67,9 @@ class StyleTransferViewController: UIViewController, ExampleViewController {
         }
     }
 
-    func loadTestImage(device: MTLDevice, commandBuffer: MTLCommandBuffer) -> MPSImage{
+    func loadTestImage(commandBuffer: MTLCommandBuffer) -> MPSImage{
         // INPUT IMAGE
-        let textureLoader = MTKTextureLoader(device: device)
+        let textureLoader = MTKTextureLoader(device: Device.shared)
         let inputTexture = try! textureLoader.newTexture(withContentsOf: Bundle.main.url(forResource: "wall-e", withExtension: "png")!, options: [MTKTextureLoaderOptionSRGB : NSNumber(value: false)])
         return MPSImage(texture: inputTexture, featureChannels: 3)
     }

Example/Example/Tests/BenderTestRunner.swift

@@ -10,12 +10,6 @@ import MetalKit
 
 public class BenderTest {
 
-    var device: MTLDevice!
-
-    init() {
-        device = MTLCreateSystemDefaultDevice()
-    }
-
     func run(completion: @escaping (Void) -> ()) {
         completion()
     }

Example/Example/Tests/ConcatTest.swift

@@ -33,7 +33,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 3, count: texture1Depth), [Float].init(repeating: 4, count: texture1Depth),
                 [Float].init(repeating: 5, count: texture1Depth), [Float].init(repeating: 6, count: texture1Depth),
             ],
-            size: LayerSize(f: texture1Depth, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: texture1Depth)
         )
         let texture2Depth = depth
         let texture2 = Texture(
@@ -48,7 +48,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 7, count: texture2Depth), [Float].init(repeating: 8, count: texture2Depth),
                 [Float].init(repeating: 8, count: texture2Depth),
             ],
-            size: LayerSize(f: texture2Depth, w: 3, h: 3)
+            size: LayerSize(h: 3, w: 3, f: texture2Depth)
         )
         let expectedDepth = depth
         let expectedData: [[Float]] = [
@@ -68,7 +68,7 @@ struct ConcatDataSet {
 
         let expected: Texture = Texture(
             data: expectedData,
-            size: LayerSize(f: expectedDepth, w: 5, h: 3)
+            size: LayerSize(h: 3, w: 5, f: expectedDepth)
         )
         return (inputTextures: [texture1, texture2], axis: .w, expected: expected)
     }
@@ -81,7 +81,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 3, count: texture1Depth), [Float].init(repeating: 4, count: texture1Depth),
                 [Float].init(repeating: 5, count: texture1Depth), [Float].init(repeating: 6, count: texture1Depth),
             ],
-            size: LayerSize(f: texture1Depth, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: texture1Depth)
         )
         let texture2Depth = depth
         let texture2 = Texture(
@@ -92,7 +92,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 9, count: texture1Depth), [Float].init(repeating: 7, count: texture1Depth),
                 [Float].init(repeating: 3, count: texture1Depth), [Float].init(repeating: 3, count: texture1Depth),
             ],
-            size: LayerSize(f: texture2Depth, w: 2, h: 5)
+            size: LayerSize(h: 5, w: 2, f: texture2Depth)
         )
         let expectedDepth = depth
         let expectedData: [[Float]] = [
@@ -108,7 +108,7 @@ struct ConcatDataSet {
 
         let expected: Texture = Texture(
             data: expectedData,
-            size: LayerSize(f: expectedDepth, w: 2, h: 8)
+            size: LayerSize(h: 8, w: 2, f: expectedDepth)
         )
         return (inputTextures: [texture1, texture2], axis: .h, expected: expected)
     }
@@ -121,7 +121,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 3, count: texture1Depth), [Float].init(repeating: 4, count: texture1Depth),
                 [Float].init(repeating: 5, count: texture1Depth), [Float].init(repeating: 6, count: texture1Depth),
             ],
-            size: LayerSize(f: texture1Depth, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: texture1Depth)
         )
         let texture2Depth = depth2
         let texture2 = Texture(
@@ -130,7 +130,7 @@ struct ConcatDataSet {
                 [Float].init(repeating: 9, count: texture2Depth), [Float].init(repeating: 10, count: texture2Depth),
                 [Float].init(repeating: 11, count: texture2Depth), [Float].init(repeating: 12, count: texture2Depth),
             ],
-            size: LayerSize(f: texture2Depth, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: texture2Depth)
         )
         let expectedDepth = depth1 + depth2
         let expectedData: [[Float]] = [
@@ -147,7 +147,7 @@ struct ConcatDataSet {
 
         let expected: Texture = Texture(
             data: expectedData,
-            size: LayerSize(f: expectedDepth, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: expectedDepth)
         )
         return (inputTextures: [texture1, texture2], axis: .f, expected: expected)
     }
@@ -167,15 +167,15 @@ class ConcatTest: BenderTest {
     }
 
     func test(inputTextures: [Texture], axis: LayerSizeAxis, expectedOutput: Texture, completion: @escaping (Void) -> ()) {
-        let styleNet = Network(device: device, inputSize: inputTextures[0].size, parameterLoader: SingleBinaryLoader(checkpoint: "lala"))
+        let styleNet = Network(inputSize: inputTextures[0].size)
 
         styleNet.start
             ->> inputTextures.map { Constant(outputTexture: $0) }
             ->> Concat(axis: axis)
 
         styleNet.initialize()
-        let metalTexture = inputTextures[0].metalTexture(with: device)
-        styleNet.run(inputImage: MPSImage(texture: metalTexture, featureChannels: inputTextures[0].depth), queue: device.makeCommandQueue()) { image in
+        let metalTexture = inputTextures[0].metalTexture(with: Device.shared)
+        styleNet.run(input: MPSImage(texture: metalTexture, featureChannels: inputTextures[0].depth)) { image in
             let textureFromGpu = Texture(metalTexture: image.texture, size: expectedOutput.size)
             assert(textureFromGpu.isEqual(to: expectedOutput))
             completion()

Example/Example/Tests/Helpers/TestData.swift

@@ -19,7 +19,7 @@ struct TestData {
                 [17.0, 18, 19, 20, 21, 22, 23, 24], [25, 26, 27, 28, 29, 30, 31, 32],
                 [33, 34, 35, 36, 37, 38, 39, 40], [41, 42, 43, 44, 45, 46, 47, 48]
             ],
-            size: LayerSize(f: 8, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: 8)
         ))
 
         textures.append(Texture(
@@ -28,7 +28,7 @@ struct TestData {
                 [17.0, 18, 19, 20, 21, 22], [25, 26, 27, 28, 29, 30],
                 [33, 34, 35, 36, 37, 38], [41, 42, 43, 44, 45, 46]
             ],
-            size: LayerSize(f: 6, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: 6)
         ))
 
         textures.append(Texture(
@@ -37,7 +37,7 @@ struct TestData {
                 [17.0, 18, 19, 20], [25, 26, 27, 28],
                 [33, 34, 35, 36], [41, 42, 43, 44]
             ],
-            size: LayerSize(f: 4, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: 4)
         ))
 
         textures.append(Texture(
@@ -46,7 +46,7 @@ struct TestData {
                 [17.0, 18, 19], [25, 26, 27],
                 [33, 34, 35], [41, 42, 43]
             ],
-            size: LayerSize(f: 3, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: 3)
         ))
 
         textures.append(Texture(
@@ -55,7 +55,7 @@ struct TestData {
                 [1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0], [1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0],
                 [1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0], [1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0]
             ],
-            size: LayerSize(f: 9, w: 2, h: 3)
+            size: LayerSize(h: 3, w: 2, f: 9)
         ))
 
         return textures

Example/Example/Tests/InstanceNormTests.swift

@@ -24,16 +24,16 @@ class InstanceNormTest: BenderTest {
     }
 
     func test(texture: Texture, completion: @escaping (Void) -> ()) {
-        let styleNet = Network(device: device, inputSize: texture.size, parameterLoader: SingleBinaryLoader(checkpoint: "lala"))
+        let styleNet = Network(inputSize: texture.size)
         let weights = [Float].init(repeating: Float(arc4random()) / Float(UINT32_MAX), count: texture.depth)
         let bias = [Float].init(repeating: Float(arc4random()) / Float(UINT32_MAX), count: texture.depth)
         let scale = Data.init(bytes: weights, count: texture.totalCount * MemoryLayout<Float>.stride)
         let shift = Data.init(bytes: bias, count: texture.totalCount * MemoryLayout<Float>.stride)
         styleNet.start ->> InstanceNorm(scale: scale, shift: shift)
         styleNet.initialize()
-        let metalTexture = texture.metalTexture(with: device)
+        let metalTexture = texture.metalTexture(with: Device.shared)
         let cpuComputed = cpuInstanceNorm(input: texture, weights: weights, bias: bias)
-        styleNet.run(inputImage: MPSImage(texture: metalTexture, featureChannels: texture.depth), queue: device.makeCommandQueue()) { image in
+        styleNet.run(input: MPSImage(texture: metalTexture, featureChannels: texture.depth)) { image in
             let textureFromGpu = Texture(metalTexture: image.texture, size: texture.size)
             assert(textureFromGpu.isEqual(to: cpuComputed, threshold: 0.002))
             completion()

Example/Example/Tests/LocalResponseNormTest.swift

@@ -32,12 +32,12 @@ class LocalResponseNormTest: BenderTest {
     }
 
     func test(texture: Texture, parameters: LocalResponseNorm.Parameters, completion: @escaping (Void) -> ()) {
-        let styleNet = Network(device: device, inputSize: texture.size, parameterLoader: SingleBinaryLoader(checkpoint: "lala"))
+        let styleNet = Network(inputSize: texture.size)
         styleNet.start ->> LocalResponseNorm(parameters: parameters, id: nil)
         styleNet.initialize()
-        let metalTexture = texture.metalTexture(with: device)
+        let metalTexture = texture.metalTexture(with: Device.shared)
         let cpuComputed = cpuLocalResponseNorm(input: texture, parameters: parameters)
-        styleNet.run(inputImage: MPSImage(texture: metalTexture, featureChannels: texture.depth), queue: device.makeCommandQueue()) { image in
+        styleNet.run(input: MPSImage(texture: metalTexture, featureChannels: texture.depth)) { image in
             let textureFromGpu = Texture(metalTexture: image.texture, size: texture.size)
             assert(textureFromGpu.isEqual(to: cpuComputed, threshold: 0.001))
             completion()

Example/Example/Tests/TextureConversionTest.swift

@@ -16,7 +16,7 @@ class TextureConversionTest: BenderTest {
     }
 
     func test(texture: Texture) {
-        let metalTexture = texture.metalTexture(with: device)
+        let metalTexture = texture.metalTexture(with: Device.shared)
         assert(Texture(metalTexture: metalTexture, size: texture.size).isEqual(to: texture))
     }