Remove Python/Numpy dependency for weights loading

Author: vvmnnnkv

FastStyleTransfer/Demo/ColabDemo.ipynb

@@ -277,9 +277,10 @@
         "let styles = [\"candy\", \"mosaic\", \"udnie\"] \n",
         "for s in styles {\n",
         "  // Download pre-trained weights\n",
-        "  downloadFile(from: \"https://github.com/vvmnnnkv/swift-models/raw/fast-style/FastStyleTransfer/Demo/weights/\\(s).npz\", to: \"\\(s).npz\")\n",
+        "  downloadFile(from: \"https://github.com/vvmnnnkv/swift-models/raw/fast-style/FastStyleTransfer/Demo/weights/\\(s).data-00000-of-00001\", to: \"\\(s).data-00000-of-00001\")\n",
+        "  downloadFile(from: \"https://github.com/vvmnnnkv/swift-models/raw/fast-style/FastStyleTransfer/Demo/weights/\\(s).index\", to: \"\\(s).index\")\n",
         "  // Load weights into model\n",
-        "  style.unsafeImport(fromNumpyArchive: \"\\(s).npz\", map: map)\n",
+        "  style.unsafeImport(fromCheckpointPath: \"\\(s)\", map: map)\n",
         "  // Apply model to image\n",
         "  let out = style(image.expandingShape(at: 0))  / 255\n",
         "  show_img(out.squeezingShape(at: 0))\n",

FastStyleTransfer/Demo/Helpers.swift

@@ -20,9 +20,9 @@ enum FileError: Error {
     case fileNotFound
 }
 
-/// Updates `model` with parameters from numpy archive file in `path`.
+/// Updates `model` with parameters from V2 checkpoint in `path`.
 func importWeights(_ model: inout TransformerNet, from path: String) throws {
-    guard FileManager.default.fileExists(atPath: path) else {
+    guard FileManager.default.fileExists(atPath: path + ".data-00000-of-00001") else {
         throw FileError.fileNotFound
     }
     // Names don't match exactly, and axes in filters need to be reversed.
@@ -74,7 +74,7 @@ func importWeights(_ model: inout TransformerNet, from path: String) throws {
         "in5.scale": ("in5.weight", nil),
         "in5.offset": ("in5.bias", nil),
     ]
-    model.unsafeImport(fromNumpyArchive: path, map: map)
+    model.unsafeImport(fromCheckpointPath: path, map: map)
 }
 
 /// Loads from `file` and returns JPEG image as HxWxC tensor of floats in (0..1) range.

FastStyleTransfer/Demo/weights/candy.npz renamed to FastStyleTransfer/Demo/weights/candy.data-00000-of-00001

@@ -1,15 +1,22 @@
 import sys
 import torch
 import numpy as np
+import tensorflow as tf
 
 # Usage:
 # python torch-convert.py model.pth model
-# (produces model.npz)
+# (produces tensorflow checkpoint model.*)
 
 if __name__ == "__main__":
-	in_file, out_file= sys.argv[1], sys.argv[2]
-	state_dict = torch.load(in_file)
-	npz = {}
-	for label, tensor in state_dict.items():
-		npz[label] = tensor.numpy()
-	np.savez(out_file, **npz)
+    in_file, out_file = sys.argv[1], sys.argv[2]
+    state_dict = torch.load(in_file)
+    variables = {}
+    tf.reset_default_graph()
+    for label, tensor in state_dict.items():
+        variables[label] = tf.get_variable(label, initializer=tensor.numpy())
+
+    init_op = tf.global_variables_initializer()
+    saver = tf.train.Saver()
+    with tf.Session() as sess:
+        sess.run(init_op)
+        save_path = saver.save(sess, out_file)

FastStyleTransfer/Demo/weights/candy.index

@@ -6,16 +6,13 @@ The model should be trainable, but so far it's only tested for inference with pr
 ## Example
 Run demo application to apply styles to jpeg images:
 ```
-swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/candy.npz --input=FastStyleTranser/Demo/examples/cat.jpg --output=candy_cat.jpg
-swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/mosaic.npz --input=FastStyleTranser/Demo/examples/cat.jpg --output=mosaic_cat.jpg
+swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/candy --input=FastStyleTranser/Demo/examples/cat.jpg --output=candy_cat.jpg
+swift run FastStyleTranserDemo --weights=FastStyleTranser/Demo/weights/mosaic --input=FastStyleTranser/Demo/examples/cat.jpg --output=mosaic_cat.jpg
 ```
 
 <img src="Demo/examples/cat.jpg" height="240" width="240" align="left">
 <img src="Demo/examples/cat_candy.jpg" height="240" width="240" align="left">
 <img src="Demo/examples/cat_mosaic.jpg" height="240" width="240">
 
-## Requirements
-Requires Python and NumPy to load weights.
-
 ## Jupyter Notebook
 Run [demo notebook](Demo/ColabDemo.ipynb) in [Colab](https://colab.research.google.com/github/vvmnnnkv/swift-models/blob/fast-style/FastStyleTransfer/Demo/ColabDemo.ipynb)!

FastStyleTransfer/Demo/weights/mosaic.npz renamed to FastStyleTransfer/Demo/weights/mosaic.data-00000-of-00001

@@ -1,6 +1,4 @@
 import TensorFlow
-import Python
-let np = Python.import("numpy")
 
 public protocol ImportableLayer: KeyPathIterable {}
 
@@ -46,14 +44,16 @@ public extension ImportableLayer {
     }
 
     /// Updates model parameters with values from `parameters`, according to `ImportMap`.
-    /// TODO add shapes check
-    mutating func unsafeImport<T>(parameters: [String: Tensor<T>], map: ImportMap) {
-        for (label, keyPath) in getRecursiveNamedKeyPaths(ofType: Tensor<T>.self) {
-            // let shape = self[keyPath: keyPath].shape
+    mutating func unsafeImport(parameters: [String: Tensor<Float>], map: ImportMap) {
+        for (label, keyPath) in getRecursiveNamedKeyPaths(ofType: Tensor<Float>.self) {
+            let shape = self[keyPath: keyPath].shape
             if let mapping = map[label], var weights = parameters[mapping.0] {
                 if let permutes = mapping.1 {
                     weights = weights.transposed(withPermutations: permutes)
                 }
+                if weights.shape != shape {
+                    fatalError("Shapes do not match for \(label): \(shape) vs. \(weights.shape)")
+                }
                 self[keyPath: keyPath] = weights
                 // print("imported \(mapping.0) \(shape) -> \(label) \(weights.shape)")
             } else if let weights = parameters[label] {
@@ -65,15 +65,16 @@ public extension ImportableLayer {
 }
 
 public extension ImportableLayer {
-    /// Updates model parameters with values from numpy archive, according to `ImportMap`.
-    mutating func unsafeImport(fromNumpyArchive file: String, map: ImportMap) {
-        let data = np.load(file)
-        var parameters = [String: Tensor<Float>]()
-        for label in data.files {
-            if let label = String(label) {
-                parameters[label] = Tensor<Float>(numpy: data[label])
-            }
-        }
+    /// Updates model parameters with values from V2 checkpoint, according to `ImportMap`.
+    mutating func unsafeImport(fromCheckpointPath path: String, map: ImportMap) {
+        let tensorNames = map.values.map { $0.0 }
+        let tensorValues = Raw.restoreV2(
+            prefix: StringTensor(path), 
+            tensorNames: StringTensor(tensorNames),
+            shapeAndSlices: StringTensor(Array(repeating: "", count: tensorNames.count)),
+            dtypes: Array(repeating: Float.tensorFlowDataType, count: tensorNames.count)
+        ).map { $0 as! Tensor<Float> }
+        let parameters = Dictionary(uniqueKeysWithValues: zip(tensorNames, tensorValues))
         unsafeImport(parameters: parameters, map: map)
     }
 }