Remove _Raw operators from models

[BradLarson]

In an effort to layer APIs and allow for runtimes other than the traditional TensorFlow eager runtime, access to Raw TensorFlow operators have been marked as _Raw. This expresses that they should be considered as internal implementation details and should not be accessed directly.

These operators are used in several places within the swift-models repository, and we should replace them with higher-level implementations of these operations or rethink how they're used. These are the operators that should be replaced:

• [x] Datasets/CIFAR10/CIFAR10.swift: _Raw.concat(concatDim:)
• [x] MiniGo/Models/PythonCheckpointReader.swift, Transformer/PythonCheckpointReader.swift: _Raw.restoreV2(prefix:)
• [x] Transformer/Model.swift: _Raw.matrixBandPart()
• [x] Transformer/Model.swift: _Raw.concatV2()
• [x] Transformer/main.swift: _Raw.multinomial(logits:, numSamples:)
• [x] Transformer/Operators.swift: _Raw.batchMatMul()
• [x] Support/Image.swift:_Raw.readFile(filename:), _Raw.decodeJpeg(contents:, channels:, dctMethod:), _Raw.reverse(), _Raw.encodeJpeg(), _Raw.writeFile(filename:, contents:), _Raw.resizeBilinear()
• [x] FastStyleTransfer/Demo/Helpers.swift: _Raw.readFile(filename:), _Raw.decodeJpeg(contents:, channels:, dctMethod:), _Raw.clipByValue(t:, clipValueMin:, clipValueMax:), _Raw.encodeJpeg(), _Raw.writeFile(filename:, contents:)
• [x] FastStyleTransfer/Utility/ImportableLayer.swift: _Raw.restoreV2(prefix:)
• [x] FastStyleTransfer/Utility/Resize.swift: _Raw.resizeNearestNeighbor(images:, size:), _Raw.resizeNearestNeighborGrad(grads:, size:)

The trickiest ones to replace will be the image operations, because that had provided cross-platform image loading and saving, and the checkpoint reading, which will require the addition of direct checkpoint parsing.

[rickwierenga]

When I tried installing this repo in Google Colab I got the following (most likely related) error:

%install-location $cwd/swift-install
%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("master"))' Datasets

error:

Installing packages:
    .package(url: "https://github.com/tensorflow/swift-models", .branch("master"))
        Datasets
With SwiftPM flags: []
Working in: /tmp/tmpr007sd63/swift-install
[1/2] Compiling ModelSupport Image.swift
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:40:26: error: use of unresolved identifier '_Raw'
        let loadedFile = _Raw.readFile(filename: StringTensor(url.absoluteString))
                         ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:41:26: error: use of unresolved identifier '_Raw'
        let loadedJpeg = _Raw.decodeJpeg(contents: loadedFile, channels: 3, dctMethod: "")
                         ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:44:23: error: use of unresolved identifier '_Raw'
                data: _Raw.reverse(loadedJpeg, dims: Tensor<Bool>([false, false, false, true])))
                      ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:62:27: error: use of unresolved identifier '_Raw'
        let encodedJpeg = _Raw.encodeJpeg(
                          ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:64:9: error: use of unresolved identifier '_Raw'
        _Raw.writeFile(filename: StringTensor(url.absoluteString), contents: encodedJpeg)
        ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:71:25: error: use of unresolved identifier '_Raw'
                tensor: _Raw.resizeBilinear(
                        ^~~~
/content/swift-install/package/.build/checkouts/swift-models/Support/Image.swift:76:25: error: use of unresolved identifier '_Raw'
                tensor: _Raw.resizeBilinear(
                        ^~~~
Install Error: swift-build returned nonzero exit code 1.

Is there any work around for now?

[BradLarson]

Colab is currently running on the 0.5 release toolchain, which is before we migrated over to the _Raw annotation here. The 0.6 release toolchain is currently under evaluation in Colab, hopefully we'll have that out soon publicly in Colab.

In the meantime, you can use the tensorflow-0.5 branch of swift-models within Colab to match up with the current toolchain. Again, it shouldn't be that much longer before Colab is updated with the 0.6 toolchain.

[Shashi456]

~@BradLarson the image based raw operations don't have underlying wrappers in Swift-apis, how would you suggest supporting these operations?~

Oh I see #273 & #274 solves most of this. So is this done?

[BradLarson]

@Shashi456 - The checkpoint loading linked above is a different issue that's currently being addressed. For this, we need platform-independent and fast image loading and saving code (JPEG and PNG as primary targets) that has no dependencies on the TensorFlow runtime or libraries. For SwiftPlot, Karthik wrapped lodepng in a Swift interface, and for harebrain / SwiftAI they wrapped libjpeg-turbo.

Ideally, I'd like to see a native-Swift implementation of PNG and JPEG loading and saving with no external dependencies. Support for acceleration within that (largely for loading) would be awesome, too. This would be useful for not just this project. I know I struggled with platform-independent image loading in GPUImage, and SwiftPlot could also benefit.

Perhaps this could be a great Summer of Code project idea, but this would be nice to have in the near term.

[Kshitij09]

@BradLarson I was in need to use _Raw.randomStandardNormal() for some testing purpose and I'm guessing I shouldn't (inferring from this issue). What alternatives do we have for such random Initializations? I found there's an init method of Tensor struct.

[BradLarson]

@Kshitij09 - You shouldn't need to use a _Raw operation for a standard normal initialization of a tensor, this initializer should do what you want.

[Kshitij09]

I'm still new to the framework and _Raw was the first thing came into my mind. Thanks! :sweat_smile:

[brettkoonce]

🥳