Skip to content
This repository was archived by the owner on Jul 1, 2023. It is now read-only.

Commit 2c96f96

Browse files
rxweirxwei
authored
Add random tensor initializers that use TensorFlow ops. (#35)
* Add `Tensor.init(randomUniform:seed:)`, `Tensor.init(randomNormal:seed:)` and `Tensor.init(glorotUniform:seed:)` that make use of TensorFlow's `StatelessRandomUniform` and `StatelessRandomNormal` operators. The original random initializers are kept so that `Tensor`'s still work well with the Swift standard library's RNG protocol and custom RNGs. * Add an initializer that uses these new random tensor initializers to `Dense` and `Conv2D`. For generality, they are probably not supposed to take a seed (but a generator + distribution instead), but we add the seed parameter today to be able to test it reliably.
1 parent ce67f8d commit 2c96f96

File tree

5 files changed

+119
-58
lines changed

5 files changed

+119
-58
lines changed

‎Sources/DeepLearning/Initializers.swift‎

Lines changed: 66 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -42,46 +42,64 @@ public extension Tensor where Scalar == Int32 {
4242
/// - shape: The dimensions of the tensor.
4343
///
4444
init(randomStandardUniform shape: TensorShape) {
45-
let dist = UniformIntegerDistribution<Scalar>()
46-
var scalars: [Scalar] = []
47-
for _ in 0 ..< shape.contiguousSize {
48-
scalars.append(dist.next(using: &PhiloxRandomNumberGenerator.global))
49-
}
50-
self.init(shape: shape, scalars: scalars)
45+
self.init(randomStandardUniform: shape, generator: &PhiloxRandomNumberGenerator.global)
5146
}
5247
}
5348

54-
public extension Tensor where Scalar: BinaryFloatingPoint,
55-
Scalar.RawSignificand: FixedWidthInteger {
49+
public extension Tensor where Scalar: BinaryFloatingPoint {
5650
/// Creates a tensor with the specified shape, randomly sampling scalar values
57-
/// from a uniform distribution between 0 and 1.
51+
/// from a uniform distribution between 0 and 1, using the default random
52+
/// number generator.
5853
///
5954
/// - Parameters:
6055
/// - shape: The dimensions of the tensor.
61-
/// - generator: Random number generator to use.
56+
/// - seed: The seed value.
57+
///
58+
init(
59+
randomUniform shape: TensorShape,
60+
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
61+
Int64.random(in: Int64.min..<Int64.max))
62+
) {
63+
self = Raw.statelessRandomUniform(
64+
shape: Tensor<Int32>((0..<shape.rank).map{shape[$0]}),
65+
seed: Tensor<Int64>([seed.0, seed.1])
66+
)
67+
}
68+
69+
/// Creates a tensor with the specified shape, randomly sampling scalar values
70+
/// from a normal distribution, using the default random number generator.
6271
///
63-
init<G: RandomNumberGenerator>(randomUniform shape: TensorShape,
64-
generator: inout G) {
65-
let dist = UniformFloatingPointDistribution<Scalar>()
66-
var scalars: [Scalar] = []
67-
for _ in 0 ..< shape.contiguousSize {
68-
scalars.append(dist.next(using: &generator))
69-
}
70-
self.init(shape: shape, scalars: scalars)
72+
/// - Parameters:
73+
/// - shape: The dimensions of the tensor.
74+
/// - seed: The seed value.
75+
///
76+
init(
77+
randomNormal shape: TensorShape,
78+
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
79+
Int64.random(in: Int64.min..<Int64.max))
80+
) {
81+
self = Raw.statelessRandomNormal(
82+
shape: Tensor<Int32>((0..<shape.rank).map{shape[$0]}),
83+
seed: Tensor<Int64>([seed.0, seed.1])
84+
)
7185
}
86+
}
7287

88+
public extension Tensor where Scalar: BinaryFloatingPoint,
89+
Scalar.RawSignificand: FixedWidthInteger {
7390
/// Creates a tensor with the specified shape, randomly sampling scalar values
74-
/// from a uniform distribution between 0 and 1, using the default random
75-
/// number generator.
91+
/// from a uniform distribution between 0 and 1.
7692
///
7793
/// - Parameters:
7894
/// - shape: The dimensions of the tensor.
95+
/// - generator: Random number generator to use.
7996
///
80-
init(randomUniform shape: TensorShape) {
97+
init<G: RandomNumberGenerator>(randomUniform shape: TensorShape,
98+
generator: inout G) {
8199
let dist = UniformFloatingPointDistribution<Scalar>()
82100
var scalars: [Scalar] = []
83101
for _ in 0 ..< shape.contiguousSize {
84-
scalars.append(dist.next(using: &PhiloxRandomNumberGenerator.global))
102+
scalars.append(dist.next(using: &generator))
85103
}
86104
self.init(shape: shape, scalars: scalars)
87105
}
@@ -106,22 +124,35 @@ public extension Tensor where Scalar: BinaryFloatingPoint,
106124
}
107125
self.init(shape: shape, scalars: scalars)
108126
}
127+
}
109128

110-
/// Creates a tensor with the specified shape, randomly sampling scalar values
111-
/// from a normal distribution, using the default random number generator.
129+
public extension Tensor where Scalar: TensorFlowFloatingPoint {
130+
private static func glorot(
131+
fromStandardUniform randomUniform: __shared Tensor<Scalar>,
132+
shape: __shared TensorShape
133+
) -> Tensor<Scalar> {
134+
let spatialDimCount = shape.count - 2
135+
let receptiveField = shape[0..<spatialDimCount].contiguousSize
136+
let fanIn = shape[shape.count - 2] * receptiveField
137+
let fanOut = shape[shape.count - 1] * receptiveField
138+
let minusOneToOne = 2 * randomUniform - 1
139+
return sqrt(Scalar(6) / Scalar(fanIn + fanOut)) * minusOneToOne
140+
}
141+
142+
/// Creates a tensor by performing Glorot uniform initialization for the specified shape,
143+
/// randomly sampling scalar values from a uniform distribution between `-limit` and `limit`,
144+
/// generated by the default random number generator, where limit is
145+
/// `sqrt(6 / (fanIn + fanOut))` and `fanIn`/`fanOut` represent the number of input and output
146+
/// features multiplied by the receptive field if present.
112147
///
113148
/// - Parameters:
114149
/// - shape: The dimensions of the tensor.
115-
/// - mean: The mean of the distribution.
116-
/// - stddev: The standard deviation of the distribution.
117150
///
118-
init(randomNormal shape: TensorShape, mean: Scalar = 0, stddev: Scalar = 1) {
119-
let dist = NormalDistribution<Scalar>(mean: mean, standardDeviation: stddev)
120-
var scalars: [Scalar] = []
121-
for _ in 0 ..< shape.contiguousSize {
122-
scalars.append(dist.next(using:&PhiloxRandomNumberGenerator.global))
123-
}
124-
self.init(shape: shape, scalars: scalars)
151+
init(glorotUniform shape: TensorShape,
152+
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
153+
Int64.random(in: Int64.min..<Int64.max))) {
154+
let uniform = Tensor(randomUniform: shape, seed: seed)
155+
self = Tensor.glorot(fromStandardUniform: uniform, shape: shape)
125156
}
126157
}
127158

@@ -137,24 +168,7 @@ public extension Tensor where Scalar: TensorFlowFloatingPoint,
137168
/// - generator: Random number generator to use.
138169
///
139170
init<G: RandomNumberGenerator>(glorotUniform shape: TensorShape, generator: inout G) {
140-
let spatialDimCount = shape.count - 2
141-
let receptiveField = shape[0..<spatialDimCount].contiguousSize
142-
let fanIn = shape[shape.count - 2] * receptiveField
143-
let fanOut = shape[shape.count - 1] * receptiveField
144-
let minusOneToOne = 2 * Tensor(randomUniform: shape, generator: &generator) - 1
145-
self = sqrt(Scalar(6) / Scalar(fanIn + fanOut)) * minusOneToOne
146-
}
147-
148-
/// Creates a tensor by performing Glorot uniform initialization for the specified shape,
149-
/// randomly sampling scalar values from a uniform distribution between `-limit` and `limit`,
150-
/// generated by the default random number generator, where limit is
151-
/// `sqrt(6 / (fanIn + fanOut))` and `fanIn`/`fanOut` represent the number of input and output
152-
/// features multiplied by the receptive field if present.
153-
///
154-
/// - Parameters:
155-
/// - shape: The dimensions of the tensor.
156-
///
157-
init(glorotUniform shape: TensorShape) {
158-
self.init(glorotUniform: shape, generator: &PhiloxRandomNumberGenerator.global)
171+
let uniform = Tensor(randomUniform: shape, generator: &generator)
172+
self = Tensor.glorot(fromStandardUniform: uniform, shape: shape)
159173
}
160174
}

‎Sources/DeepLearning/Layer.swift‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,21 @@ public extension Dense where Scalar.RawSignificand: FixedWidthInteger {
221221
}
222222
}
223223

224+
public extension Dense {
225+
init(
226+
inputSize: Int,
227+
outputSize: Int,
228+
activation: @escaping Activation = identity,
229+
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
230+
Int64.random(in: Int64.min..<Int64.max))
231+
) {
232+
self.init(weight: Tensor(glorotUniform: [Int32(inputSize), Int32(outputSize)],
233+
seed: seed),
234+
bias: Tensor(zeros: [Int32(outputSize)]),
235+
activation: activation)
236+
}
237+
}
238+
224239
@_fixed_layout
225240
public struct Conv2D<Scalar: TensorFlowFloatingPoint>: Layer {
226241
public var filter: Tensor<Scalar>
@@ -269,6 +284,27 @@ public extension Conv2D where Scalar.RawSignificand: FixedWidthInteger {
269284
}
270285
}
271286

287+
public extension Conv2D {
288+
init(
289+
filterShape: (Int, Int, Int, Int),
290+
strides: (Int, Int) = (1, 1),
291+
padding: Padding = .valid,
292+
activation: @escaping Activation = identity,
293+
seed: (Int64, Int64) = (Int64.random(in: Int64.min..<Int64.max),
294+
Int64.random(in: Int64.min..<Int64.max))
295+
) {
296+
let filterTensorShape = TensorShape([
297+
Int32(filterShape.0), Int32(filterShape.1),
298+
Int32(filterShape.2), Int32(filterShape.3)])
299+
self.init(
300+
filter: Tensor(glorotUniform: filterTensorShape, seed: seed),
301+
bias: Tensor(zeros: TensorShape([Int32(filterShape.3)])),
302+
activation: activation,
303+
strides: (Int32(strides.0), Int32(strides.1)),
304+
padding: padding)
305+
}
306+
}
307+
272308
@_fixed_layout
273309
public struct BatchNorm<Scalar: TensorFlowFloatingPoint>: Layer {
274310
/// The batch dimension.

‎Tests/DeepLearningTests/PRNGTests.swift‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,12 +235,22 @@ final class PRNGTests: XCTestCase {
235235
}
236236
}
237237

238+
func testTensorFlowInitializers() {
239+
let shape: TensorShape = [128, 128]
240+
let seed: (Int64, Int64) = (41, 42)
241+
measure {
242+
_ = Tensor<Float>(randomUniform: shape, seed: seed)
243+
_ = Tensor<Float>(randomNormal: shape, seed: seed)
244+
}
245+
}
246+
238247
static var allTests = [
239248
("testARC4", testARC4),
240249
("testUniformDistribution", testUniformDistribution),
241250
("testNormalDistribution", testNormalDistribution),
242251
("testUniformIntegerDistribution", testUniformIntegerDistribution),
243252
("testThreefry", testThreefry),
244253
("testPhilox", testPhilox),
254+
("testTensorFlowInitializers", testTensorFlowInitializers),
245255
]
246256
}

‎Tests/DeepLearningTests/SequentialTests.swift‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@ import XCTest
1818
final class SequentialTests: XCTestCase {
1919
func testSequential() {
2020
struct Model: Layer {
21-
var dense1 = Dense<Float>(inputSize: 2, outputSize: 4, activation: relu)
22-
var dense2 = Dense<Float>(inputSize: 4, outputSize: 1, activation: relu)
21+
var dense1 = Dense<Float>(inputSize: 2, outputSize: 4, activation: relu,
22+
seed: (0xfffffff, 0xfeeff))
23+
var dense2 = Dense<Float>(inputSize: 4, outputSize: 1, activation: relu,
24+
seed: (0xfeffeffe, 0xfffe))
2325

2426
@differentiable
2527
func applied(to input: Tensor<Float>, in context: Context) -> Tensor<Float> {

‎Tests/DeepLearningTests/TrivialModelTests.swift‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,20 +18,19 @@ import XCTest
1818
final class TrivialModelTests: XCTestCase {
1919
func testXOR() {
2020
struct Classifier: Layer {
21-
static var generator = PhiloxRandomNumberGenerator(uint64Seed: 51243124)
2221
var l1, l2: Dense<Float>
2322
init(hiddenSize: Int) {
2423
l1 = Dense<Float>(
2524
inputSize: 2,
2625
outputSize: hiddenSize,
2726
activation: relu,
28-
generator: &Classifier.generator
27+
seed: (0xfffffff, 0xfeeff)
2928
)
3029
l2 = Dense<Float>(
3130
inputSize: hiddenSize,
3231
outputSize: 1,
3332
activation: relu,
34-
generator: &Classifier.generator
33+
seed: (0xfeffeffe, 0xfffe)
3534
)
3635
}
3736
@differentiable
@@ -46,7 +45,7 @@ final class TrivialModelTests: XCTestCase {
4645
let y: Tensor<Float> = [[0], [1], [1], [0]]
4746

4847
let trainingContext = Context(learningPhase: .training)
49-
for _ in 0..<2000 {
48+
for _ in 0..<3000 {
5049
let 𝛁model = classifier.gradient { classifier -> Tensor<Float> in
5150
let ŷ = classifier.applied(to: x, in: trainingContext)
5251
return meanSquaredError(predicted: ŷ, expected: y)

0 commit comments

Comments
 (0)