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

Commit 8264ac0

Browse files
tanmayb123rxwei
tanmayb123
authored and
rxwei
committed
Add operations from standard library (#63)
* Add operations from standard library * Op formatting * Style changes * Style changes * Formatting * Formatting * Fix pooling ops * Format pooling ops
1 parent 905f0fb commit 8264ac0

File tree

1 file changed

+325
-0
lines changed

1 file changed

+325
-0
lines changed

‎Sources/DeepLearning/Operators.swift

Lines changed: 325 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,328 @@ public func round<Scalar: BinaryFloatingPoint>(_ x: Tensor<Scalar>) -> Tensor<Sc
2626
public func identity<Scalar>(_ x: Tensor<Scalar>) -> Tensor<Scalar> {
2727
return x
2828
}
29+
30+
//===------------------------------------------------------------------------------------------===//
31+
// Normalization
32+
//===------------------------------------------------------------------------------------------===//
33+
34+
public extension Tensor where Scalar: TensorFlowFloatingPoint {
35+
// TODO: Verify that these calculations are correct.
36+
@inlinable
37+
func _vjpBatchNormalized(
38+
alongAxis axis: Int32,
39+
offset: Tensor,
40+
scale: Tensor,
41+
epsilon: Scalar
42+
) -> (Tensor, (Tensor) -> (Tensor, Tensor, Tensor)) {
43+
let value = batchNormalized(alongAxis: axis, offset: offset, scale: scale,
44+
epsilon: epsilon)
45+
return (value, { v in
46+
let mean = self.mean(alongAxes: axis)
47+
let squaredDiff: Tensor = Raw.squaredDifference(self, mean)
48+
let variance = squaredDiff.mean(alongAxes: axis)
49+
50+
let diff = self - mean
51+
let inv = rsqrt(variance + epsilon)
52+
let norm = diff * inv
53+
54+
let dNorm = v * scale
55+
let dVariance = -(dNorm * diff).sum(alongAxes: axis) / 2 * pow(inv, -3)
56+
let dMean = (-dNorm * inv).sum(alongAxes: axis) +
57+
dVariance * (-diff * 2).mean(alongAxes: axis)
58+
let dOffset = v.sum(alongAxes: axis)
59+
let dScale = (norm * v).sum(alongAxes: axis)
60+
let dim = Tensor(Tensor<Int32>(self.shapeTensor[axis]))
61+
let tmp = (dNorm * inv) + (dVariance * 2 * dMean / dim)
62+
let dSelf = tmp + (dMean / dim)
63+
return (dSelf, dOffset, dScale)
64+
})
65+
}
66+
}
67+
68+
public extension Tensor where Scalar: BinaryFloatingPoint {
69+
/// Computes the batch normalized tensor along the specified axis.
70+
///
71+
/// Specifically, returns `(self - mu)/(var + epsilon) * gamma + beta` where
72+
/// `mu` and `var` are respectively the mean and variance of `self` along
73+
/// `axis`.
74+
///
75+
/// - Parameters:
76+
/// - axis: The batch dimension.
77+
/// - offset: The offset, also known as beta.
78+
/// - scale: The scale, also known as gamma.
79+
/// - epsilon: A small value added to the denominator for numerical
80+
/// stability.
81+
@inlinable
82+
@differentiable(
83+
wrt: (self, offset, scale), vjp: _vjpBatchNormalized
84+
where Scalar : TensorFlowFloatingPoint
85+
)
86+
func batchNormalized(
87+
alongAxis axis: Int32,
88+
offset: Tensor = Tensor(0),
89+
scale: Tensor = Tensor(1),
90+
epsilon: Scalar = 0.001
91+
) -> Tensor {
92+
let mean = self.mean(alongAxes: axis)
93+
let squaredDiff: Tensor = Raw.squaredDifference(self, mean)
94+
let variance = squaredDiff.mean(alongAxes: axis)
95+
let inv = rsqrt(variance + epsilon) * scale
96+
return self * inv + offset - mean * inv
97+
}
98+
}
99+
100+
//===------------------------------------------------------------------------------------------===//
101+
// Convolution and pooling
102+
//===------------------------------------------------------------------------------------------===//
103+
104+
/// A padding scheme. Used by padding, convolution, and pooling ops.
105+
// @_frozen // SR-9739
106+
public enum Padding {
107+
/// The "valid" padding scheme.
108+
case valid
109+
/// The "same" padding scheme.
110+
case same
111+
}
112+
113+
public extension Padding {
114+
@inlinable
115+
var raw: Raw.Padding {
116+
switch self {
117+
case .same: return .same
118+
case .valid: return .valid
119+
}
120+
}
121+
}
122+
123+
extension Tensor where Scalar: TensorFlowFloatingPoint {
124+
/// TensorFlow builtin conv2d gradient helper for the input.
125+
@inlinable
126+
@differentiable(
127+
wrt: (filter, backpropOutput),
128+
vjp: _vjpTFConv2DBackpropInput(_:_:_:_:_:)
129+
)
130+
func _TFConv2DBackpropInput(
131+
shape: Tensor<Int32>,
132+
filter: Tensor,
133+
backpropOutput: Tensor,
134+
strides: (Int32, Int32, Int32, Int32),
135+
padding: Padding
136+
) -> Tensor {
137+
return Raw.conv2DBackpropInput(
138+
inputSizes: shape,
139+
filter: filter,
140+
outBackprop: backpropOutput,
141+
strides: [strides.0, strides.1, strides.2, strides.3],
142+
padding: padding.raw)
143+
}
144+
145+
/// TensorFlow builtin conv2d gradient helper for the filter.
146+
@inlinable
147+
@differentiable(
148+
wrt: (input, backpropOutput),
149+
vjp: _vjpTFConv2DBackpropFilter(_:_:_:_:_:)
150+
)
151+
func _TFConv2DBackpropFilter(
152+
input: Tensor,
153+
filterSizes: Tensor<Int32>,
154+
backpropOutput: Tensor,
155+
strides: (Int32, Int32, Int32, Int32),
156+
padding: Padding
157+
) -> Tensor {
158+
return Raw.conv2DBackpropFilter(
159+
input,
160+
filterSizes: filterSizes,
161+
outBackprop: backpropOutput,
162+
strides: [strides.0, strides.1, strides.2, strides.3],
163+
padding: padding.raw)
164+
}
165+
166+
@inlinable
167+
func _vjpTFConv2DBackpropInput(
168+
_ shape: Tensor<Int32>,
169+
_ filter: Tensor,
170+
_ backpropOutput: Tensor,
171+
_ strides: (Int32, Int32, Int32, Int32),
172+
_ padding: Padding
173+
) -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
174+
let value = _TFConv2DBackpropInput(shape: shape, filter: filter,
175+
backpropOutput: backpropOutput,
176+
strides: strides, padding: padding)
177+
return (value, { v in
178+
return (
179+
self._TFConv2DBackpropFilter(input: v, filterSizes: shape,
180+
backpropOutput: backpropOutput,
181+
strides: strides, padding: padding),
182+
v.convolved2D(withFilter: filter, strides: strides, padding: padding)
183+
)
184+
})
185+
}
186+
187+
@inlinable
188+
func _vjpTFConv2DBackpropFilter(
189+
_ input: Tensor,
190+
_ filterSizes: Tensor<Int32>,
191+
_ backpropOutput: Tensor,
192+
_ strides: (Int32, Int32, Int32, Int32),
193+
_ padding: Padding
194+
) -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
195+
let value = _TFConv2DBackpropFilter(input: input, filterSizes: filterSizes,
196+
backpropOutput: backpropOutput,
197+
strides: strides, padding: padding)
198+
return (value, { v in
199+
return (
200+
self._TFConv2DBackpropInput(shape: filterSizes, filter: v,
201+
backpropOutput: backpropOutput,
202+
strides: strides, padding: padding),
203+
input.convolved2D(withFilter: v, strides: strides, padding: padding)
204+
)
205+
})
206+
}
207+
208+
@inlinable
209+
func _vjpConvolved2D(
210+
filter: Tensor,
211+
strides: (Int32, Int32, Int32, Int32),
212+
padding: Padding
213+
) -> (Tensor, (Tensor) -> (Tensor, Tensor)) {
214+
let value = convolved2D(withFilter: filter, strides: strides,
215+
padding: padding)
216+
return (value, { v in
217+
return (
218+
self._TFConv2DBackpropInput(
219+
shape: self.shapeTensor, filter: filter, backpropOutput: v,
220+
strides: strides, padding: padding
221+
),
222+
self._TFConv2DBackpropFilter(
223+
input: self, filterSizes: filter.shapeTensor, backpropOutput: v,
224+
strides: strides, padding: padding
225+
)
226+
)
227+
})
228+
}
229+
230+
@inlinable
231+
func _vjpMaxPooled(
232+
kernelSize: (Int32, Int32, Int32, Int32),
233+
strides: (Int32, Int32, Int32, Int32),
234+
padding: Padding
235+
) -> (Tensor, (Tensor) -> Tensor) {
236+
// TODO: Currently this is not higher order differentiable. Redefine in
237+
// closed form.
238+
let value = maxPooled(kernelSize: kernelSize, strides: strides,
239+
padding: padding)
240+
return (value, { v in
241+
return Raw.maxPoolGradV2(
242+
origInput: self,
243+
origOutput: value,
244+
grad: v,
245+
ksize: Tensor<Int32>([kernelSize.0, kernelSize.1, kernelSize.2, kernelSize.3]),
246+
strides: Tensor<Int32>([strides.0, strides.1, strides.2, strides.3]),
247+
padding: padding.raw
248+
)
249+
})
250+
}
251+
252+
@inlinable
253+
func _vjpAveragePooled(
254+
kernelSize: (Int32, Int32, Int32, Int32),
255+
strides: (Int32, Int32, Int32, Int32),
256+
padding: Padding
257+
) -> (Tensor, (Tensor) -> Tensor) {
258+
// TODO: Currently this is not higher order differentiable. Redefine in
259+
// closed form.
260+
let value = averagePooled(kernelSize: kernelSize, strides: strides,
261+
padding: padding)
262+
return (value, { v in
263+
return Raw.avgPoolGrad(
264+
origInputShape: self.shapeTensor,
265+
grad: v,
266+
ksize: [kernelSize.0, kernelSize.1, kernelSize.2, kernelSize.3],
267+
strides: [strides.0, strides.1, strides.2, strides.3],
268+
padding: padding.raw
269+
)
270+
})
271+
}
272+
}
273+
274+
public extension Tensor where Scalar: FloatingPoint {
275+
/// Computes a 2-D convolution using `self` as input, with the specified
276+
/// filter, strides, and padding.
277+
///
278+
/// - Parameters:
279+
/// - filter: The convolution filter.
280+
/// - strides: The strides of the sliding filter for each dimension of the
281+
/// input.
282+
/// - padding: The padding for the operation.
283+
/// - Precondition: `self` must have rank 4.
284+
/// - Precondition: `filter` must have rank 4.
285+
@inlinable @inline(__always)
286+
@differentiable(
287+
wrt: (self, filter), vjp: _vjpConvolved2D(filter:strides:padding:)
288+
where Scalar : TensorFlowFloatingPoint
289+
)
290+
func convolved2D(
291+
withFilter filter: Tensor,
292+
strides: (Int32, Int32, Int32, Int32),
293+
padding: Padding
294+
) -> Tensor {
295+
return Raw.conv2D(
296+
self,
297+
filter: filter,
298+
strides: [strides.0, strides.1, strides.2, strides.3],
299+
padding: padding.raw)
300+
}
301+
302+
/// Computes a 2-D max pooling, with the specified kernel sizes, strides, and
303+
/// padding.
304+
///
305+
/// - Parameters:
306+
/// - kernelSize: The dimensions of the pooling kernel.
307+
/// - strides: The strides of the sliding filter for each dimension of the
308+
/// input.
309+
/// - padding: The padding for the operation.
310+
@inlinable @inline(__always)
311+
@differentiable(
312+
wrt: self, vjp: _vjpMaxPooled(kernelSize:strides:padding:)
313+
where Scalar : TensorFlowFloatingPoint
314+
)
315+
func maxPooled(
316+
kernelSize: (Int32, Int32, Int32, Int32),
317+
strides: (Int32, Int32, Int32, Int32),
318+
padding: Padding
319+
) -> Tensor {
320+
return Raw.maxPoolV2(
321+
self,
322+
ksize: Tensor<Int32>([kernelSize.0, kernelSize.1,
323+
kernelSize.2, kernelSize.3]),
324+
strides: Tensor<Int32>([strides.0, strides.1,
325+
strides.2, strides.3]),
326+
padding: padding.raw)
327+
}
328+
329+
/// Computes a 2-D average pooling, with the specified kernel sizes, strides,
330+
/// and padding.
331+
///
332+
/// - Parameters:
333+
/// - kernelSize: The dimensions of the pooling kernel.
334+
/// - strides: The strides of the sliding filter for each dimension of the
335+
/// input.
336+
/// - padding: The padding for the operation.
337+
@inlinable @inline(__always)
338+
@differentiable(
339+
wrt: self, vjp: _vjpAveragePooled(kernelSize:strides:padding:)
340+
where Scalar : TensorFlowFloatingPoint
341+
)
342+
func averagePooled(
343+
kernelSize: (Int32, Int32, Int32, Int32),
344+
strides: (Int32, Int32, Int32, Int32),
345+
padding: Padding
346+
) -> Tensor {
347+
return Raw.avgPool(
348+
value: self,
349+
ksize: [kernelSize.0, kernelSize.1, kernelSize.2, kernelSize.3],
350+
strides: [strides.0, strides.1, strides.2, strides.3],
351+
padding: padding.raw)
352+
}
353+
}

0 commit comments

Comments
 (0)