Skip to content
This repository was archived by the owner on Apr 23, 2025. It is now read-only.

Commit a08917a

Browse files
shadajshadaj
committed
Add all layers needed for LeNet-5
1 parent 2a8ca19 commit a08917a

File tree

7 files changed

+182
-62
lines changed

7 files changed

+182
-62
lines changed

‎Examples/LeNet-MNIST/main.swift

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,18 +22,16 @@ let batchSize = 128
2222
let dataset = MNIST(batchSize: batchSize)
2323
// The LeNet-5 model, equivalent to `LeNet` in `ImageClassificationModels`.
2424

25-
let denseDef = AutoDenseDefinition<Float>(outputSize: 120, activation: relu)
26-
.then(AutoDenseDefinition(outputSize: 84, activation: relu))
27-
.then(AutoDenseDefinition(outputSize: 10))
25+
let denseDef = AutoConv2D<Float>(filterShape: (5, 5), outputChannels: 6, padding: .same, activation: relu)
26+
.then(AutoAvgPool2D(poolSize: (2, 2), strides: (2, 2)))
27+
.then(AutoConv2D(filterShape: (5, 5), outputChannels: 16, activation: relu))
28+
.then(AutoAvgPool2D(poolSize: (2, 2), strides: (2, 2)))
29+
.then(AutoFlatten())
30+
.then(AutoDense(outputSize: 120, activation: relu))
31+
.then(AutoDense(outputSize: 84, activation: relu))
32+
.then(AutoDense(outputSize: 10))
2833

29-
var classifier = Sequential {
30-
Conv2D<Float>(filterShape: (5, 5, 1, 6), padding: .same, activation: relu)
31-
AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
32-
Conv2D<Float>(filterShape: (5, 5, 6, 16), activation: relu)
33-
AvgPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
34-
Flatten<Float>()
35-
denseDef.buildModel(inputShape: 400)
36-
}
34+
var classifier = denseDef.buildModel(inputShape: (28, 28, 1))
3735

3836
let optimizer = SGD(for: classifier, learningRate: 0.1)
3937

‎Models/LayerInit/AutoConv.swift

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
import TensorFlow
2+
3+
public struct AutoConv2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
4+
let filterShape: (Int, Int)
5+
let outputChannels: Int
6+
let strides: (Int, Int)
7+
let padding: Padding
8+
let dilations: (Int, Int)
9+
let activation: Conv2D<Scalar>.Activation
10+
let useBias: Bool
11+
let filterInitializer: ParameterInitializer<Scalar>
12+
let biasInitializer: ParameterInitializer<Scalar>
13+
14+
public typealias InstanceType = Conv2D<Scalar>
15+
public typealias InputShape = (Int, Int, Int)
16+
public typealias OutputShape = (Int, Int, Int)
17+
18+
public init(
19+
filterShape: (Int, Int),
20+
outputChannels: Int,
21+
strides: (Int, Int) = (1, 1),
22+
padding: Padding = .valid,
23+
dilations: (Int, Int) = (1, 1),
24+
activation: @escaping Conv2D<Scalar>.Activation = identity,
25+
useBias: Bool = true,
26+
filterInitializer: @escaping ParameterInitializer<Scalar> = glorotUniform(),
27+
biasInitializer: @escaping ParameterInitializer<Scalar> = zeros()
28+
) {
29+
self.filterShape = filterShape
30+
self.outputChannels = outputChannels
31+
self.strides = strides
32+
self.padding = padding
33+
self.dilations = dilations
34+
self.activation = activation
35+
self.useBias = useBias
36+
self.filterInitializer = filterInitializer
37+
self.biasInitializer = biasInitializer
38+
}
39+
40+
public func buildModelWithOutputShape(inputShape: (Int, Int, Int)) -> (InstanceType, (Int, Int, Int)) {
41+
let outputShape: (Int, Int, Int)
42+
if (padding == .valid) {
43+
outputShape = (
44+
Int(ceil(Float(inputShape.0 - filterShape.0 + 1) / Float(strides.0))),
45+
Int(ceil(Float(inputShape.1 - filterShape.1 + 1) / Float(strides.1))),
46+
outputChannels
47+
)
48+
} else {
49+
outputShape = (
50+
Int(ceil(Float(inputShape.0) / Float(strides.0))),
51+
Int(ceil(Float(inputShape.1) / Float(strides.1))),
52+
outputChannels
53+
)
54+
}
55+
56+
return (Conv2D<Scalar>(
57+
filterShape: (filterShape.0, filterShape.1, inputShape.2, outputChannels),
58+
strides: strides, padding: padding, dilations: dilations,
59+
activation: activation, useBias: useBias,
60+
filterInitializer: filterInitializer, biasInitializer: biasInitializer
61+
), outputShape)
62+
}
63+
}

‎Models/LayerInit/AutoDense.swift

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import TensorFlow
2+
3+
public struct AutoDense<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
4+
let outputSize: Int;
5+
let activation: Dense<Scalar>.Activation
6+
7+
public typealias InstanceType = Dense<Scalar>
8+
public typealias InputShape = Int
9+
public typealias OutputShape = Int
10+
11+
public init(outputSize: Int, activation: @escaping Dense<Scalar>.Activation = identity) {
12+
self.outputSize = outputSize
13+
self.activation = activation
14+
}
15+
16+
public func buildModelWithOutputShape(inputShape: Int) -> (InstanceType, Int) {
17+
return (Dense<Scalar>(inputSize: inputShape, outputSize: self.outputSize, activation: self.activation), self.outputSize)
18+
}
19+
}

‎Models/LayerInit/AutoFlatten.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import TensorFlow
2+
3+
public struct AutoFlatten<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
4+
public typealias InstanceType = Flatten<Scalar>
5+
public typealias InputShape = (Int, Int, Int)
6+
public typealias OutputShape = Int
7+
8+
public init() {}
9+
10+
public func buildModelWithOutputShape(inputShape: (Int, Int, Int)) -> (InstanceType, Int) {
11+
return (Flatten<Scalar>(), inputShape.0 * inputShape.1 * inputShape.2)
12+
}
13+
}

‎Models/LayerInit/AutoLayer.swift

Lines changed: 5 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -2,60 +2,14 @@ import TensorFlow
22

33
public protocol AutoLayer {
44
associatedtype InstanceType: Layer
5-
func buildModelWithOutputShape(inputShape: Int) -> (InstanceType, Int)
5+
associatedtype InputShape
6+
associatedtype OutputShape
7+
8+
func buildModelWithOutputShape(inputShape: InputShape) -> (InstanceType, OutputShape)
69
}
710

811
extension AutoLayer {
9-
public func buildModel(inputShape: Int) -> InstanceType {
12+
public func buildModel(inputShape: InputShape) -> InstanceType {
1013
return self.buildModelWithOutputShape(inputShape: inputShape).0
1114
}
1215
}
13-
14-
public struct AutoSequencedDefinition<Layer1: AutoLayer, Layer2: AutoLayer>: AutoLayer
15-
where
16-
Layer1.InstanceType.Output == Layer2.InstanceType.Input,
17-
Layer1.InstanceType.TangentVector.VectorSpaceScalar == Layer2.InstanceType.TangentVector.VectorSpaceScalar {
18-
let first: Layer1
19-
let second: Layer2
20-
21-
public typealias InstanceType = Sequential<Layer1.InstanceType, Layer2.InstanceType>
22-
23-
public init(first: Layer1, second: Layer2) {
24-
self.first = first
25-
self.second = second
26-
}
27-
28-
public func buildModelWithOutputShape(inputShape: Int) -> (InstanceType, Int) {
29-
let (firstInstance, firstOutputShape) = first.buildModelWithOutputShape(inputShape: inputShape)
30-
let (secondInstance, secondOutputShape) = second.buildModelWithOutputShape(inputShape: firstOutputShape)
31-
return (Sequential(firstInstance, secondInstance), secondOutputShape)
32-
}
33-
}
34-
35-
extension AutoSequencedDefinition {
36-
public func then<T: AutoLayer>(_ other: T) -> AutoSequencedDefinition<AutoSequencedDefinition<Layer1, Layer2>, T> {
37-
return AutoSequencedDefinition<AutoSequencedDefinition<Layer1, Layer2>, T>(first: self, second: other)
38-
}
39-
}
40-
41-
public struct AutoDenseDefinition<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
42-
let outputSize: Int;
43-
let activation: Dense<Scalar>.Activation
44-
45-
public typealias InstanceType = Dense<Scalar>
46-
47-
public init(outputSize: Int, activation: @escaping Dense<Scalar>.Activation = identity) {
48-
self.outputSize = outputSize
49-
self.activation = activation
50-
}
51-
52-
public func buildModelWithOutputShape(inputShape: Int) -> (InstanceType, Int) {
53-
return (Dense<Scalar>(inputSize: inputShape, outputSize: self.outputSize, activation: self.activation), self.outputSize)
54-
}
55-
}
56-
57-
extension AutoDenseDefinition {
58-
public func then<T: AutoLayer>(_ other: T) -> AutoSequencedDefinition<AutoDenseDefinition<Scalar>, T> {
59-
return AutoSequencedDefinition<AutoDenseDefinition<Scalar>, T>(first: self, second: other)
60-
}
61-
}

‎Models/LayerInit/AutoPool.swift

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import TensorFlow
2+
3+
public struct AutoAvgPool2D<Scalar>: AutoLayer where Scalar: TensorFlowFloatingPoint {
4+
let poolSize: (Int, Int)
5+
let strides: (Int, Int)
6+
let padding: Padding
7+
8+
public typealias InstanceType = AvgPool2D<Scalar>
9+
public typealias InputShape = (Int, Int, Int)
10+
public typealias OutputShape = (Int, Int, Int)
11+
12+
public init(
13+
poolSize: (Int, Int),
14+
strides: (Int, Int) = (1, 1),
15+
padding: Padding = .valid
16+
) {
17+
self.poolSize = poolSize
18+
self.strides = strides
19+
self.padding = padding
20+
}
21+
22+
public func buildModelWithOutputShape(inputShape: (Int, Int, Int)) -> (InstanceType, (Int, Int, Int)) {
23+
let outputShape: (Int, Int, Int)
24+
if (padding == .valid) {
25+
outputShape = (
26+
Int(ceil(Float(inputShape.0 - poolSize.0 + 1) / Float(strides.0))),
27+
Int(ceil(Float(inputShape.1 - poolSize.1 + 1) / Float(strides.1))),
28+
inputShape.2
29+
)
30+
} else {
31+
outputShape = (
32+
Int(ceil(Float(inputShape.0) / Float(strides.0))),
33+
Int(ceil(Float(inputShape.1) / Float(strides.1))),
34+
inputShape.2
35+
)
36+
}
37+
38+
return (AvgPool2D<Scalar>(
39+
poolSize: poolSize,
40+
strides: strides,
41+
padding: padding
42+
), outputShape)
43+
}
44+
}

‎Models/LayerInit/AutoSequenced.swift

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import TensorFlow
2+
3+
public struct AutoSequencedDefinition<Layer1: AutoLayer, Layer2: AutoLayer>: AutoLayer
4+
where
5+
Layer1.OutputShape == Layer2.InputShape,
6+
Layer1.InstanceType.Output == Layer2.InstanceType.Input,
7+
Layer1.InstanceType.TangentVector.VectorSpaceScalar == Layer2.InstanceType.TangentVector.VectorSpaceScalar {
8+
let first: Layer1
9+
let second: Layer2
10+
11+
public typealias InstanceType = Sequential<Layer1.InstanceType, Layer2.InstanceType>
12+
13+
public init(first: Layer1, second: Layer2) {
14+
self.first = first
15+
self.second = second
16+
}
17+
18+
public func buildModelWithOutputShape(inputShape: Layer1.InputShape) -> (InstanceType, Layer2.OutputShape) {
19+
let (firstInstance, firstOutputShape) = first.buildModelWithOutputShape(inputShape: inputShape)
20+
let (secondInstance, secondOutputShape) = second.buildModelWithOutputShape(inputShape: firstOutputShape)
21+
return (Sequential(firstInstance, secondInstance), secondOutputShape)
22+
}
23+
}
24+
25+
extension AutoLayer {
26+
public func then<T: AutoLayer>(_ other: T) -> AutoSequencedDefinition<Self, T> {
27+
return AutoSequencedDefinition<Self, T>(first: self, second: other)
28+
}
29+
}

0 commit comments

Comments
 (0)