%install '.package(path: "$cwd/FastaiNotebook_06_cuda")' FastaiNotebook_06_cuda
Installing packages: .package(path: "/home/jekbradbury/git/fastai_docs/dev_swift/FastaiNotebook_06_cuda") FastaiNotebook_06_cuda With SwiftPM flags: [] Working in: /tmp/tmpp95uqama/swift-install /home/jekbradbury/swift/usr/bin/swift-build: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/lib/swift/linux/libFoundation.so) /home/jekbradbury/swift/usr/bin/swift-build: /home/jekbradbury/anaconda3/lib/libcurl.so.4: no version information available (required by /home/jekbradbury/swift/usr/lib/swift/linux/libFoundation.so) Fetching https://github.com/mxcl/Path.swift Fetching https://github.com/JustHTTP/Just Fetching https://github.com/latenitesoft/NotebookExport Completed resolution in 3.65s Cloning https://github.com/latenitesoft/NotebookExport Resolving https://github.com/latenitesoft/NotebookExport at 0.5.0 Cloning https://github.com/JustHTTP/Just Resolving https://github.com/JustHTTP/Just at 0.7.1 Cloning https://github.com/mxcl/Path.swift Resolving https://github.com/mxcl/Path.swift at 0.16.2 /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) Compile Swift Module 'Just' (1 sources) Compile Swift Module 'Path' (9 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'NotebookExport' (2 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_00_load_data' (1 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_01_matmul' (2 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_01a_fastai_layers' (3 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_02_fully_connected' (4 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_02a_why_sqrt5' (5 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_03_minibatch_training' (6 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_04_callbacks' (7 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_05_anneal' (8 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_05b_early_stopping' (9 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'FastaiNotebook_06_cuda' (10 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Compile Swift Module 'jupyterInstalledPackages' (1 sources) /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) /home/jekbradbury/swift/usr/bin/swift: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift) Linking ./.build/x86_64-unknown-linux/debug/libjupyterInstalledPackages.so /home/jekbradbury/swift/usr/bin/swiftc: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swiftc) /home/jekbradbury/swift/usr/bin/swift-autolink-extract: /home/jekbradbury/anaconda3/lib/libuuid.so.1: no version information available (required by /home/jekbradbury/swift/usr/bin/swift-autolink-extract) Initializing Swift... Installation complete!
import FastaiNotebook_06_cuda
%include "EnableIPythonDisplay.swift"
IPythonDisplay.shell.enable_matplotlib("inline")
('inline', 'module://ipykernel.pylab.backend_inline')
//export
import Path
import TensorFlow
import Python
Let's start by building our own batchnorm layer from scratch. Eventually we want something like this to work:
class AlmostBatchNorm<Scalar: TensorFlowFloatingPoint> { // : Layer
// Configuration hyperparameters
let momentum, epsilon: Scalar
// Running statistics
var runningMean, runningVariance: Tensor<Scalar>
// Trainable parameters
var scale, offset: Tensor<Scalar>
init(featureCount: Int, momentum: Scalar = 0.9, epsilon: Scalar = 1e-5) {
(self.momentum, self.epsilon) = (momentum, epsilon)
(scale, offset) = (Tensor(ones: [featureCount]), Tensor(zeros: [featureCount]))
(runningMean, runningVariance) = (Tensor(0), Tensor(1))
}
func call(_ input: Tensor<Scalar>) -> Tensor<Scalar> {
let mean, variance: Tensor<Scalar>
switch Context.local.learningPhase {
case .training:
mean = input.mean(alongAxes: [0, 1, 2])
variance = input.variance(alongAxes: [0, 1, 2])
runningMean += (mean - runningMean) * (1 - momentum)
runningVariance += (variance - runningVariance) * (1 - momentum)
case .inference:
(mean, variance) = (runningMean, runningVariance)
}
let normalizer = rsqrt(variance + epsilon) * scale
return (input - mean) * normalizer + offset
}
}
But there are some automatic differentiation limitations (lack of support for classes and control flow) that make this impossible for now, so we'll need a few workarounds. A Reference will let us update running statistics without making the layer a class or declaring the applied method mutating:
//export
public class Reference<T> {
public var value: T
public init(_ value: T) { self.value = value }
}
The following snippet will let us differentiate a layer's applied method if it's composed of training and inference implementations that are each differentiable:
//export
public protocol LearningPhaseDependent: FALayer {
associatedtype Input
associatedtype Output
@differentiable func forwardTraining(to input: Input) -> Output
@differentiable func forwardInference(to input: Input) -> Output
}
extension LearningPhaseDependent {
public func forward(_ input: Input) -> Output {
switch Context.local.learningPhase {
case .training: return forwardTraining(to: input)
case .inference: return forwardInference(to: input)
}
}
@differentiating(forward)
func gradForward(_ input: Input) ->
(value: Output, pullback: (Self.Output.CotangentVector) ->
(Self.CotangentVector, Self.Input.CotangentVector)) {
switch Context.local.learningPhase {
case .training: return valueWithPullback(at: input) { $0.forwardTraining(to: $1) }
case .inference: return valueWithPullback(at: input) { $0.forwardInference(to: $1) }
}
}
}
Now we can implement a BatchNorm that we can use in our models:
//export
public protocol Norm: FALayer where Input == Tensor<Float>, Output == Tensor<Float> {
init(featureCount: Int, epsilon: Float)
}
public struct FABatchNorm: LearningPhaseDependent, Norm {
// Configuration hyperparameters
@noDerivative var momentum, epsilon: Float
// Running statistics
@noDerivative let runningMean, runningVariance: Reference<Tensor<Float>>
// Trainable parameters
public var scale, offset: Tensor<Float>
public init(featureCount: Int, momentum: Float, epsilon: Float = 1e-5) {
self.momentum = momentum
self.epsilon = epsilon
self.scale = Tensor(ones: [featureCount])
self.offset = Tensor(zeros: [featureCount])
self.runningMean = Reference(Tensor(0))
self.runningVariance = Reference(Tensor(1))
}
public init(featureCount: Int, epsilon: Float = 1e-5) {
self.init(featureCount: featureCount, momentum: 0.9, epsilon: epsilon)
}
@differentiable
public func forwardTraining(to input: Tensor<Float>) -> Tensor<Float> {
let mean = input.mean(alongAxes: [0, 1, 2])
let variance = input.variance(alongAxes: [0, 1, 2])
runningMean.value += (mean - runningMean.value) * (1 - momentum)
runningVariance.value += (variance - runningVariance.value) * (1 - momentum)
let normalizer = rsqrt(variance + epsilon) * scale
return (input - mean) * normalizer + offset
}
@differentiable
public func forwardInference(to input: Tensor<Float>) -> Tensor<Float> {
let (mean, variance) = (runningMean.value, runningVariance.value)
let normalizer = rsqrt(variance + epsilon) * scale
return (input - mean) * normalizer + offset
}
// Things that should probably be synthesized/inherited, but aren't
public typealias Input = Tensor<Float>
public typealias Output = Tensor<Float>
}
//export
public struct ConvNorm<NormType: Norm & FALayer>: FALayer
where NormType.AllDifferentiableVariables == NormType.CotangentVector {
public var conv: FANoBiasConv2D<Float>
public var norm: NormType
public init(_ cIn: Int, _ cOut: Int, ks: Int = 3, stride: Int = 2){
self.conv = FANoBiasConv2D(
filterShape: (ks, ks, cIn, cOut),
strides: (stride, stride),
padding: .same,
activation: relu)
self.norm = NormType(featureCount: cOut, epsilon: 1e-5)
}
@differentiable
public func forward(_ input: Tensor<Float>) -> Tensor<Float> {
return norm(conv(input))
}
}
//export
public struct CnnModelNormed<NormType: Norm & FALayer>: FALayer
where NormType.AllDifferentiableVariables == NormType.CotangentVector {
public var convs: [ConvNorm<NormType>]
public var pool = FAGlobalAvgPool2D<Float>()
public var flatten = Flatten<Float>()
public var linear: FADense<Float>
public init(channelIn: Int, nOut: Int, filters: [Int]){
convs = []
let allFilters = [channelIn] + filters
for i in 0..<filters.count { convs.append(ConvNorm<NormType>(allFilters[i], allFilters[i+1])) }
linear = FADense<Float>(filters.last!, nOut)
}
@differentiable
public func forward(_ input: TF) -> TF {
return input.sequenced(through: convs, pool, flatten, linear)
}
}
Let's benchmark this batchnorm implementation!
func benchmark(forward: () -> (), backward: () -> ()) {
print("forward:")
time(repeating: 10, forward)
print("backward:")
time(repeating: 10, backward)
}
let input = Tensor<Float>(randomUniform: [64, 28, 28, 32])
let norm = FABatchNorm(featureCount: 32)
let pb = pullback(at: input) { x in norm(x) }
benchmark(forward: { norm(input) }, backward: { pb(input) })
Yikes, that's pretty bad. Luckily, TensorFlow has a built-in fused batchnorm layer. Let's see how the performance looks for that:
let input = Tensor<Float>(randomUniform: [64, 28, 28, 32])
let norm = FABatchNorm(featureCount: 32)
let bnresult = Raw.fusedBatchNormV2(
input, scale: norm.scale, offset: norm.offset,
mean: Tensor<Float>([] as [Float]), variance: Tensor<Float>([] as [Float]),
epsilon: Double(norm.epsilon))
benchmark(
forward: {
Raw.fusedBatchNormV2(
input, scale: norm.scale, offset: norm.offset,
mean: Tensor<Float>([] as [Float]), variance: Tensor<Float>([] as [Float]),
epsilon: Double(norm.epsilon))
},
backward: {
Raw.fusedBatchNormGradV2(
yBackprop: input, input, scale: Tensor<Float>(norm.scale),
reserveSpace1: bnresult.reserveSpace1,
reserveSpace2: bnresult.reserveSpace2,
epsilon: Double(norm.epsilon))
})
struct PullbackArgs<T : TensorGroup, U : TensorGroup> : TensorGroup {
let input: T
let cotangent: U
}
class CompiledFunction<Input: Differentiable & TensorGroup, Output: Differentiable & TensorGroup> {
let f: @differentiable (Input) -> Output
init(_ f: @escaping @differentiable (Input) -> Output) {
self.f = f
}
}
func xlaCompiled<T : Differentiable & TensorGroup, U : Differentiable & TensorGroup>(
_ fn: @escaping @differentiable (T) -> U) -> CompiledFunction<T, U>
where T.CotangentVector : TensorGroup, U.CotangentVector : TensorGroup {
let xlaCompiledFn: (T) -> U = _graph(fn, useXLA: true)
let xlaCompiledPullback = _graph(
{ (pbArgs: PullbackArgs<T, U.CotangentVector>) in
pullback(at: pbArgs.input, in: fn)(pbArgs.cotangent) },
useXLA: true
)
return CompiledFunction(differentiableFunction { x in
(value: xlaCompiledFn(x), pullback: { v in
xlaCompiledPullback(PullbackArgs(input: x, cotangent: v))})
})
}
struct TrainingKernelInput: TensorGroup, Differentiable, AdditiveArithmetic {
let input: Tensor<Float>
let scale: Tensor<Float>
let offset: Tensor<Float>
let runningMean: Tensor<Float>
let runningVariance: Tensor<Float>
let momentum: Tensor<Float>
let epsilon: Tensor<Float>
}
struct TrainingKernelOutput: TensorGroup, Differentiable, AdditiveArithmetic {
let normalized: Tensor<Float>
let newRunningMean: Tensor<Float>
let newRunningVariance: Tensor<Float>
}
@differentiable
func trainingKernel(_ input: TrainingKernelInput) -> TrainingKernelOutput {
let mean = input.input.mean(alongAxes: [0, 1, 2])
let variance = input.input.variance(alongAxes: [0, 1, 2])
let invMomentum = Tensor<Float>(1) - input.momentum
let newRunningMean = input.runningMean * input.momentum + mean * invMomentum
let newRunningVariance = input.runningVariance * input.momentum + variance * invMomentum
let normalizer = rsqrt(variance + input.epsilon) * input.scale
let normalized = (input.input - mean) * normalizer + input.offset
return TrainingKernelOutput(
normalized: normalized,
newRunningMean: newRunningMean,
newRunningVariance: newRunningVariance
)
}
let input = Tensor<Float>(randomUniform: [64, 28, 28, 32])
let norm = FABatchNorm(featureCount: 32)
let compiledTrainingKernel = xlaCompiled(trainingKernel)
let kernelInput = TrainingKernelInput(
input: input,
scale: norm.scale,
offset: norm.offset,
runningMean: norm.runningMean.value,
runningVariance: norm.runningVariance.value,
momentum: Tensor(norm.momentum),
epsilon: Tensor(norm.epsilon))
let pb = pullback(at: kernelInput) { x in compiledTrainingKernel.f(x) }
let kernelOutput = compiledTrainingKernel.f(kernelInput)
benchmark(
forward: { compiledTrainingKernel.f(kernelInput) },
backward: { pb(kernelOutput) })