Bender/Example/BatchNormTest.swift at master · xmartlabs/Bender · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
//
//  BatchNormTest.swift
//  Example
//
//  Created by Mathias Claassen on 2/13/18.
//

import Accelerate
import AVFoundation
import MetalBender
import MetalKit
import MetalPerformanceShaders
import MetalPerformanceShadersProxy

struct BatchNormDataSet {

    static let epsilon: Float = 0.001
    static let mean = [Float](repeating: 0.5, count: 8)
    static let variance = [Float](repeating: 2, count: 8)
    static let scale = [Float](repeating: 1.5, count: 8)
    static let offset = [Float](repeating: 1, count: 8)

    static func resultWithScale(val: Float) -> Float {
        return (val - mean[0]) * scale[0] / sqrt(variance[0] + epsilon) + offset[0]
    }

    static func resultSimple(val: Float) -> Float {
        return (val - mean[0]) / sqrt(variance[0] + epsilon)
    }

    typealias Test = (inputTexture: Texture, useScale: Bool, expected: Texture)

    static let testData: [Test] = [
        test1(useScale: true, depth: 8),
        test1(useScale: false, depth: 8),
        test2(useScale: true, depth: 4),
        test2(useScale: false, depth: 4),
    ]

    static func test1(useScale: Bool, depth: Int = 8) -> Test {
        let texture = Texture(
            data: [[Float].init(repeating: 1, count: depth), [Float].init(repeating: 2, count: depth),
                   [Float].init(repeating: 3, count: depth), [Float].init(repeating: 4, count: depth),
                   [Float].init(repeating: 5, count: depth), [Float].init(repeating: 6, count: depth)],
            size: LayerSize(h: 3, w: 2, f: depth)
        )

        let expected: Texture = Texture(
            data: texture.map(op: useScale ? resultWithScale : resultSimple),
            size: LayerSize(h: 3, w: 2, f: depth)
        )
        return (inputTexture: texture, useScale: useScale, expected: expected)
    }

    static func test2(useScale: Bool, depth: Int = 8) -> Test {
        var data = [[Float]]()
        for _ in 0..<3 {
            for _ in 0..<3 {
                data.append([Float].random(count: depth))
            }
        }
        let texture = Texture(data: data, size: LayerSize(h: 3, w: 3, f: depth))

        let expected: Texture = Texture(data: texture.map(op: useScale ? resultWithScale : resultSimple),
                                        size: LayerSize(h: 3, w: 3, f: depth))
        return (inputTexture: texture, useScale: useScale, expected: expected)
    }

}

class BatchNormTest: BenderTest {

    override func run(completion: @escaping () -> ()) {
        var tests: [CompletionSerializer.CompletableFunction] = []
        BatchNormDataSet.testData.forEach { testData in
            tests.append( { completion in
                self.test(inputTexture: testData.inputTexture,
                          useScale: testData.useScale,
                          expectedOutput: testData.expected,
                          completion: completion)
            })
        }
        CompletionSerializer(completableFunctions: tests).run(completion: completion)
    }

    func test(inputTexture: Texture, useScale: Bool, expectedOutput: Texture, completion: @escaping () -> ()) {
        let styleNet = Network(inputSize: inputTexture.size)

        styleNet.start
            ->> BatchNorm(mean: BatchNormDataSet.mean.toData(count: inputTexture.size.f),
                          variance: BatchNormDataSet.variance.toData(count: inputTexture.size.f),
                          offset: useScale ? BatchNormDataSet.offset.toData(count: inputTexture.size.f) : nil,
                          scale: useScale ? BatchNormDataSet.scale.toData(count: inputTexture.size.f) : nil,
                          epsilon: BatchNormDataSet.epsilon)

        styleNet.initialize()
        let metalTexture = inputTexture.metalTexture(with: Device.shared)
        styleNet.run(input: MPSImage(texture: metalTexture, featureChannels: inputTexture.depth)) { image in
            let textureFromGpu = Texture(metalTexture: image!.texture, size: expectedOutput.size)
            assert(textureFromGpu.isEqual(to: expectedOutput, threshold: 0.01))
            completion()
        }
    }

}