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) }