1、设置输入:
let input = Input()
或者
let input = Input(width: 100, height: 100, channels: 3)
2、创建网络:
let output = input --> Resize(width: 28, height: 28) --> Convolution(kernel: (5, 5), channels: 20, activation: relu, name: "conv1") --> Dense(neurons: 10, name: "dense1") --> Softmax()
3、链接网络、加载参数
model = Model(input: input, output: output)
let success = model.compile(device: device, inflightBuffers: 3) { name, count, type in return ParameterLoaderBundle(name: name, count: count, suffix: type == .weights ? "_W" : "_b", ext: "bin") } if success { print(model.summary()) }
4、预测阶段:
model.encode(commandBuffer: commandBuffer, texture: inputTexture, inflightIndex: i)
let probabilities = model.outputImage(inflightIndex: i).toFloatArray() let top5 = probabilities.top(k: 5) let top5Labels = top5.map { x -> (String, Float) in (labels[x.0], x.1) }