diff --git a/Models/ImageClassification/SqueezeNet.swift b/Models/ImageClassification/SqueezeNet.swift new file mode 100644 index 00000000000..18901da8a51 --- /dev/null +++ b/Models/ImageClassification/SqueezeNet.swift @@ -0,0 +1,116 @@ +// Copyright 2019 The TensorFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import TensorFlow + +// Original Paper: +// SqueezeNet: AlexNet Level Accuracy with 50X Fewer Parameters +// Forrest N. Iandola, Song Han, Matthew W. Moskewicz, Khalid Ashraf, William J. Dally +// and Kurt Keutzer +// https://arxiv.org/pdf/1602.07360.pdf + +public struct Fire: Layer { + public var squeeze: Conv2D + public var expand1: Conv2D + public var expand3: Conv2D + + public init( + inputFilterCount: Int, + squeezeFilterCount: Int, + expand1FilterCount: Int, + expand3FilterCount: Int + ) { + squeeze = Conv2D( + filterShape: (1, 1, inputFilterCount, squeezeFilterCount), + activation: relu) + expand1 = Conv2D( + filterShape: (1, 1, squeezeFilterCount, expand1FilterCount), + activation: relu) + expand3 = Conv2D( + filterShape: (3, 3, squeezeFilterCount, expand3FilterCount), + padding: .same, + activation: relu) + } + + @differentiable + public func callAsFunction(_ input: Tensor) -> Tensor { + let squeezed = squeeze(input) + let expanded1 = expand1(squeezed) + let expanded3 = expand3(squeezed) + return expanded1.concatenated(with: expanded3, alongAxis: -1) + } +} + +public struct SqueezeNet: Layer { + public var conv1 = Conv2D(filterShape: (7, 7, 3, 96), strides: (2, 2), padding: .same) + public var maxPool1 = MaxPool2D(poolSize: (3, 3), strides: (2, 2)) + public var fire2 = Fire( + inputFilterCount: 96, + squeezeFilterCount: 16, + expand1FilterCount: 64, + expand3FilterCount: 64) + public var fire3 = Fire( + inputFilterCount: 128, + squeezeFilterCount: 16, + expand1FilterCount: 64, + expand3FilterCount: 64) + public var fire4 = Fire( + inputFilterCount: 128, + squeezeFilterCount: 32, + expand1FilterCount: 128, + expand3FilterCount: 128) + public var maxPool4 = MaxPool2D(poolSize: (3, 3), strides: (2, 2)) + public var fire5 = Fire( + inputFilterCount: 256, + squeezeFilterCount: 32, + expand1FilterCount: 128, + expand3FilterCount: 128) + public var fire6 = Fire( + inputFilterCount: 256, + squeezeFilterCount: 48, + expand1FilterCount: 192, + expand3FilterCount: 192) + public var fire7 = Fire( + inputFilterCount: 384, + squeezeFilterCount: 48, + expand1FilterCount: 192, + expand3FilterCount: 192) + public var fire8 = Fire( + inputFilterCount: 384, + squeezeFilterCount: 64, + expand1FilterCount: 256, + expand3FilterCount: 256) + public var maxPool8 = MaxPool2D(poolSize: (3, 3), strides: (2, 2)) + public var fire9 = Fire( + inputFilterCount: 512, + squeezeFilterCount: 64, + expand1FilterCount: 256, + expand3FilterCount: 256) + public var conv10: Conv2D + public var avgPool10 = AvgPool2D(poolSize: (13, 13), strides: (1, 1)) + public var dropout = Dropout(probability: 0.5) + + public init(classCount: Int) { + conv10 = Conv2D(filterShape: (1, 1, 512, classCount), strides: (1, 1)) + } + + @differentiable + public func callAsFunction(_ input: Tensor) -> Tensor { + let convolved1 = input.sequenced(through: conv1, maxPool1) + let fired1 = convolved1.sequenced(through: fire2, fire3, fire4, maxPool4, fire5, fire6) + let fired2 = fired1.sequenced(through: fire7, fire8, maxPool8, fire9) + let convolved2 = fired2.sequenced(through: dropout, conv10, avgPool10) + return convolved2 + } +}