A CNN Mnist Model¶

In [1]:

%install '.package(path: "$cwd/FastaiNotebook_05b_early_stopping")' FastaiNotebook_05b_early_stopping

Installing packages:
	.package(path: "/home/ubuntu/fastai_docs/dev_swift/FastaiNotebook_05b_early_stopping")
		FastaiNotebook_05b_early_stopping
With SwiftPM flags: []
Working in: /tmp/tmpyihp3nn8/swift-install
Fetching https://github.com/mxcl/Path.swift
Fetching https://github.com/JustHTTP/Just
Completed resolution in 3.15s
Cloning https://github.com/mxcl/Path.swift
Resolving https://github.com/mxcl/Path.swift at 0.16.2
Cloning https://github.com/JustHTTP/Just
Resolving https://github.com/JustHTTP/Just at 0.7.1
Compile Swift Module 'Just' (1 sources)
Compile Swift Module 'Path' (9 sources)
Compile Swift Module 'FastaiNotebook_05b_early_stopping' (10 sources)
Compile Swift Module 'jupyterInstalledPackages' (1 sources)
Linking ./.build/x86_64-unknown-linux/debug/libjupyterInstalledPackages.so
Initializing Swift...
Installation complete!

Load data¶

In [2]:

import FastaiNotebook_05b_early_stopping
%include "EnableIPythonDisplay.swift"
IPythonDisplay.shell.enable_matplotlib("inline")

Out[2]:

('inline', 'module://ipykernel.pylab.backend_inline')

In [3]:

// export
import Path
import TensorFlow
import Python

In [4]:

let plt = Python.import("matplotlib.pyplot")

In [5]:

let data = mnistDataBunch(flat: false, bs: 512)

In [6]:

let firstBatch = data.train.first(where: { _ in true })!
let batchShape = firstBatch.xb.shape
let batchSize = batchShape.dimensions[0]
let exampleSideSize = batchShape.dimensions[1]
assert(exampleSideSize == batchShape.dimensions[2])
print("Batch size: \(batchSize)")
print("Example side size: \(exampleSideSize)")

let classCount = firstBatch.yb.shape.dimensions[0]
print("Class count: \(classCount)")

Batch size: 512
Example side size: 28
Class count: 512

In [7]:

firstBatch.xb.shape

Out[7]:

▿ TensorShape
  ▿ dimensions : 3 elements
    - 0 : 512
    - 1 : 28
    - 2 : 28

In [8]:

// export
extension Learner {
    public class AddChannel: Delegate {
        public override func batchWillStart(learner: Learner) {
            learner.currentInput = learner.currentInput!.expandingShape(at: -1)
        }
    }
    
    public func makeAddChannel() -> AddChannel { return AddChannel() }
}

In [9]:

//export 
public func conv<Scalar>(_ cIn: Int, _ cOut: Int, ks: Int = 3, stride: Int = 2) -> FAConv2D<Scalar> {
    return FAConv2D<Scalar>(filterShape: (ks, ks, cIn, cOut), 
                            strides: (stride,stride), 
                            padding: .same, 
                            activation: relu)
}

public struct CnnModel: Layer {
    public var convs: [FAConv2D<Float>]
    public var pool = FAAdaptiveAvgPool2D<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(conv(allFilters[i], allFilters[i+1])) }
        linear = FADense<Float>(inputSize: filters.last!, outputSize: nOut)
    }
    
    @differentiable
    public func applied(to input: TF) -> TF {
        return input.sequenced(through: convs, pool, flatten, linear)
    }
}

In [10]:

let model = CnnModel(channelIn: 1, nOut: 10, filters: [8, 16, 32, 32])

In [11]:

// Test that data goes through the model as expected.
let predictions = model.applied(to: firstBatch.xb.expandingShape(at: -1))
print(predictions.shape)
print(predictions[0])

TensorShape(dimensions: [512, 10])
[-0.0056363903, -0.0055124904,   0.003962202, 0.00024041766,  0.0010704098,   0.012630586,
   -0.01890578,  -0.012497755, -0.0075372434,  -0.001092382]

Compare training on CPU and GPU¶

In [12]:

let opt = SimpleSGD<CnnModel, Float>(learningRate: 0.4)
func modelInit() -> CnnModel { return CnnModel(channelIn: 1, nOut: 10, filters: [8, 16, 32, 32]) }
let learner = Learner(data: data, lossFunction: softmaxCrossEntropy, optimizer: opt, initializingWith: modelInit)
let recorder = learner.makeDefaultDelegates(metrics: [accuracy])
learner.addDelegates([learner.makeNormalize(mean: mnistStats.mean, std: mnistStats.std),
                      learner.makeAddChannel()])

In [13]:

// This happens on the GPU (if you have one and it's configured correctly).
// I tried this on a GCE 8vCPU 30GB + Tesla P100:
// - time: ~4.3s
// - nvidia-smi shows ~10% GPU-Util while this is running
time {
    try! learner.fit(1)
}

Epoch 0: [0.3500186, 0.8953]                                                    
4725.451923 ms

In [14]:

// This happens on the CPU.
// I tried this on a GCE 8vCPU 30GB + Tesla P100:
// - time: ~6.3s
// - nvidia-smi shows 0% GPU-Util while this is running
time {
    withDevice(.cpu) {
        try! learner.fit(1)
    }
}

Epoch 0: [0.1639871, 0.9526]                                                    
6769.385988 ms

Collect Layer Activation Statistics¶

In [15]:

class ActivationStatistics: LayerDelegate<Tensor<Float>> {
    var activationMeans: [Float] = []
    var activationStds: [Float] = []    
    override func didProduceActivation(_ activation: Tensor<Float>) {
        activationMeans.append(activation.mean().scalar!)
        activationStds.append(activation.standardDeviation().reshaped(to: []).scalar!)
    }
}

In [16]:

extension KeyPathIterable {    
    mutating func initializeLayerDelegates<T>(with initializer: () -> LayerDelegate<T>) {
        for kp in recursivelyAllWritableKeyPaths(to: LayerDelegate<T>.self) {
            self[keyPath: kp] = initializer()
        }
    }
    
    func layerDelegates<T, D: LayerDelegate<T>>(havingType: D.Type) -> [D] {
        var result: [D] = []
        for kp in recursivelyAllWritableKeyPaths(to: LayerDelegate<T>.self) {
            guard let d = self[keyPath: kp] as? D else { continue }
            result.append(d)
        }
        return result
    }
}

In [17]:

let learner = Learner(data: data, lossFunction: softmaxCrossEntropy, optimizer: opt, initializingWith: modelInit)
let recorder = learner.makeDefaultDelegates(metrics: [accuracy])
learner.addDelegates([learner.makeNormalize(mean: mnistStats.mean, std: mnistStats.std),
                      learner.makeAddChannel()])

learner.model.initializeLayerDelegates(with: { ActivationStatistics() })

In [18]:

// This LayerDelegate stuff slows it down to ~6s/epoch.
time {
    try! learner.fit(2)
}

Epoch 0: [0.3566188, 0.8989]                                                    
Epoch 1: [0.16688569, 0.9504]                                                   
9983.915673 ms

In [19]:

let activationStatistics = learner.model.layerDelegates(havingType: ActivationStatistics.self)
for stats in activationStatistics {
    plt.plot(stats.activationMeans)
}
plt.legend(Array(1...activationStatistics.count))
plt.show()

Out[19]:

None

In [20]:

for stats in activationStatistics {
    plt.plot(stats.activationStds)
}
plt.legend(Array(1...activationStatistics.count))
plt.show()

Out[20]:

None

Export¶

In [18]:

notebookToScript(fname: (Path.cwd / "06_cuda.ipynb").string)

In [ ]: