 TensorFlow.orgで表示 |  GoogleColabで実行 |  GitHubでソースを表示 |
このチュートリアルでは、独自のカスタムデリバティブを定義し、デリバティブ手術を実行し、わずか5行のSwiftで独自の勾配チェックポインティングAPIを実装する方法を示します。
カスタムデリバティブの宣言
微分可能なパラメーターと結果を持つ任意のSwift関数のカスタム導関数を定義できます。そうすることで、C関数をインポートして微分可能にすることもできます。
import Glibc
func sillyExp(_ x: Float) -> Float {
let 𝑒 = Float(M_E)
print("Taking 𝑒(\(𝑒)) to the power of \(x)!")
return pow(𝑒, x)
}
@derivative(of: sillyExp)
func sillyDerivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
let y = sillyExp(x)
return (value: y, pullback: { v in v * y })
}
print("exp(3) =", sillyExp(3))
print("𝛁exp(3) =", gradient(of: sillyExp)(3))
Taking 𝑒(2.7182817) to the power of 3.0! exp(3) = 20.085535 Taking 𝑒(2.7182817) to the power of 3.0! 𝛁exp(3) = 20.085535
デリバティブの伝播を停止します
機械学習のユースケースでは一般に「停止勾配」として知られているメソッドwithoutDerivative(at:)は、導関数の伝播を停止します。
さらに、 withoutDerivative(at:)は、Swiftコンパイラーが区別しないものを識別し、より効率的な派生物を生成するのに役立つ場合があります。関数の導関数が常にゼロであることが検出できる場合、Swiftコンパイラーは警告を生成します。 withoutDerivative(at:)明示的に使用すると、その警告は無音になります。
let x: Float = 2.0
let y: Float = 3.0
let xyGradient = gradient(at: x, y) { x, y in
sin(sin(sin(x))) + withoutDerivative(at: cos(cos(cos(y))))
}
print(xyGradient)
(-0.18009877, 0.0)
派生手術
メソッドwithDerivative(_:)は、任意の操作（ミューwithDerivative(_:)を含む）を、囲んでいる関数のバックプロパゲーション中の値で勾配上で実行します。
これを使用して、バックプロパゲーションのデバッグまたは実験的な調整を行います。
それはどこでも動作します
標準ライブラリによって提供されるすべての微分APIは、 Differentiableプロトコルに準拠するすべての型（ Float 、 Double 、 Float80 、SIMDベクトル、さらには独自の型）に対して一般的に定義されています。
Differentiableプロトコルの詳細については、技術文書「微分可能タイプ」をお読みください。
var x: Float = 30
let xGradient = gradient(at: x) { x -> Float in
// Print the partial derivative with respect to the result of `sin(x)`.
let a = sin(x).withDerivative { print("∂+/∂sin = \($0)") }
// Force the partial derivative with respect to `x` to be `0.5`.
let b = log(x.withDerivative { (dx: inout Float) in
print("∂log/∂x = \(dx), but rewritten to 0.5");
dx = 0.5
})
return a + b
}
print(xGradient)
∂log/∂x = 0.033333335, but rewritten to 0.5 ∂+/∂sin = 1.0 0.65425146
ニューラルネットワークモジュールで使用する
単純なFloat関数で使用したのと同じように、 Swift for TensorFlow Deep Learning Libraryを使用して構築された次のニューラルネットワークなど、任意の数値アプリケーションで使用できます。
import TensorFlow
struct MLP: Layer {
var layer1 = Dense<Float>(inputSize: 2, outputSize: 10, activation: relu)
var layer2 = Dense<Float>(inputSize: 10, outputSize: 1, activation: relu)
@differentiable
func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
let h0 = layer1(input).withDerivative { print("∂L/∂layer1 =", $0) }
return layer2(h0)
}
}
var classifier = MLP()
let optimizer = SGD(for: classifier, learningRate: 0.02)
let x: Tensor<Float> = [[0, 0], [0, 1], [1, 0], [1, 1]]
let y: Tensor<Float> = [0, 1, 1, 0]
for _ in 0..<10 {
let 𝛁model = gradient(at: classifier) { classifier -> Tensor<Float> in
let ŷ = classifier(x).withDerivative { print("∂L/∂ŷ =", $0) }
let loss = (ŷ - y).squared().mean()
print("Loss: \(loss)")
return loss
}
optimizer.update(&classifier, along: 𝛁model)
}
Loss: 0.4562089
∂L/∂ŷ = [[ -0.25],
[-0.13678059],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.08600927, -0.061670735, -0.07495588, -0.012159594, 0.061164618, -0.09354251,
0.0068358597, 0.048978966, -0.033759538, 0.044154454],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.45417115
∂L/∂ŷ = [[-0.24798042],
[ -0.1329543],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.15593305, -0.11180778, -0.13589348, -0.022429364, 0.1108902, -0.16987386,
0.012393274, 0.088590585, -0.061205354, 0.080051124],
[ 0.083603255, -0.05994556, -0.07285907, -0.012025467, 0.059453603, -0.091077596,
0.006644634, 0.047497697, -0.03281515, 0.042919282],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.44949773
∂L/∂ŷ = [[ -0.2423476],
[-0.12556008],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.15239106, -0.10926809, -0.13280669, -0.022285527, 0.10837136, -0.16628888,
0.0121117635, 0.08638263, -0.059815086, 0.07823278],
[ 0.078953676, -0.0566117, -0.06880703, -0.011546112, 0.056147106, -0.086154126,
0.006275094, 0.044754762, -0.030990142, 0.040532336],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.44514722
∂L/∂ŷ = [[-0.23689735],
[-0.11843509],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.14896388, -0.10681072, -0.12981996, -0.022123232, 0.10593416, -0.16281205,
0.011839378, 0.084260456, -0.05846988, 0.07647338],
[ 0.0744734, -0.053399235, -0.06490254, -0.011060348, 0.052961003, -0.08139669,
0.00591901, 0.042125396, -0.02923159, 0.038232304],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.4410977
∂L/∂ŷ = [[ -0.23162389],
[-0.111571014],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.14564787, -0.10443306, -0.12693009, -0.021944579, 0.103576005, -0.15943976,
0.011575826, 0.082220234, -0.057168312, 0.07477104],
[ 0.070157185, -0.050304405, -0.061141014, -0.010570494, 0.049891572, -0.07680061,
0.005575965, 0.0396047, -0.02753743, 0.036016494],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.4373285
∂L/∂ŷ = [[ -0.22652149],
[-0.104959674],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.14243942, -0.10213252, -0.124133974, -0.021751495, 0.101294346, -0.15616852,
0.011320825, 0.08025828, -0.055908963, 0.073123924],
[ 0.065999895, -0.04732353, -0.057517994, -0.010078645, 0.046935156, -0.072361335,
0.0052455515, 0.037188005, -0.02590565, 0.03388227],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.43382046
∂L/∂ŷ = [[ -0.22158477],
[-0.098593265],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.13933516, -0.09990668, -0.121428646, -0.021545762, 0.09908678, -0.15299515,
0.011074103, 0.07837118, -0.054690503, 0.07153028],
[ 0.061996624, -0.044453084, -0.054029197, -0.009586701, 0.04408827, -0.06807459,
0.0049273786, 0.034870945, -0.024334323, 0.03182712],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.43055546
∂L/∂ŷ = [[-0.21680838],
[-0.09246406],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.1363317, -0.09775314, -0.11881118, -0.021329017, 0.0969509, -0.14991656,
0.010835394, 0.07655567, -0.053511616, 0.06998841],
[ 0.058142506, -0.04168959, -0.050670385, -0.009096362, 0.04134745, -0.06393615,
0.00462106, 0.032649327, -0.022821542, 0.029848535],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.42751682
∂L/∂ŷ = [[-0.21218723],
[-0.08656461],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.13342586, -0.09566958, -0.11627879, -0.021102766, 0.09488445, -0.14692979,
0.010604444, 0.07480866, -0.052371047, 0.068496644],
[ 0.05443286, -0.039029684, -0.04743748, -0.008609154, 0.038709376, -0.059941966,
0.0043262243, 0.030519187, -0.021365467, 0.02794412],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
Loss: 0.42468888
∂L/∂ŷ = [[-0.20771626],
[-0.08088773],
[ -0.25],
[ -0.25]]
∂L/∂layer1 = [[ 0.13061446, -0.09365374, -0.11382869, -0.020868381, 0.09288515, -0.14403199,
0.010380999, 0.0731272, -0.051267542, 0.06705336],
[ 0.05086317, -0.036470126, -0.044326544, -0.008126452, 0.036170825, -0.05608815,
0.004042512, 0.028476795, -0.019964326, 0.026111552],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0],
[ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0]]
メモリを節約するためのバックプロパゲーション中のアクティベーションの再計算（チェックポイント）
チェックポインティングは、メモリを節約するためのリバースモード自動微分の従来の手法です。導関数を計算するための元の計算で大きな中間値を保存するのではなく、中間値はバックプロパゲーション中に必要に応じて再計算されます。
この手法は、最新のディープラーニングライブラリでも実現されています。 Swiftでは、API withRecomputationInPullbacks(_:)使用すると、バックプロパゲーション中に何を再計算するかを制御でき、すべてのDifferentiable可能タイプで使用できます。
しかし、今日は、ほんの数行のコードで、独自のグラデーションチェックポインティングAPIを最初から定義する方法を学びましょう。
グラデーションチェックポイントAPI
標準ライブラリ関数differentiableFunction(from:)観点から、独自の勾配チェックポイントAPI makeRecomputedInGradient(_:)を定義できます。これは、微分関数（「ベクトル-ヤコビ行列」とも呼ばれます）から直接微分可能関数を作成するための省略形です。 （VJP）関数」）。
前に見たように、微分関数は元の関数の結果のタプルとプルバッククロージャを返します。私たちは、返しoriginal(x)中value: 、およびコールpullback(at:in:)上original 、再び本来の機能を評価し、プルバックを取得します。
/// Given a differentiable function, returns the same differentiable function except when
/// derivatives of this function are being computed. In that case, values in the original function needed
/// for computing the derivatives will be recomputed, instead of being captured by the differential or pullback.
///
/// - Parameter body: The body of the differentiable function.
/// - Returns: The same differentiable function whose derivatives, when computed, will recompute
/// some values from the original function.
func makeRecomputedInGradient<T: Differentiable, U: Differentiable>(
_ original: @escaping @differentiable (T) -> U
) -> @differentiable (T) -> U {
return differentiableFunction { x in
(value: original(x), pullback: { v in pullback(at: x, in: original)(v) })
}
}
動作することを確認します
let input: Float = 10.0
print("Running original computation...")
// Differentiable multiplication with checkpointing.
let square = makeRecomputedInGradient { (x: Float) -> Float in
print(" Computing square...")
return x * x
}
// Differentiate `f(x) = (cos(x))^2`.
let (output, backprop) = valueWithPullback(at: input) { input -> Float in
return square(cos(input))
}
print("Running backpropagation...")
let grad = backprop(1)
print("Gradient = \(grad)")
Running original computation... Computing square... Running backpropagation... Computing square... Gradient = -0.9129453
それをニューラルネットワークモジュールに拡張する
この例では、単純な畳み込みニューラルネットワークを定義します。
struct Model: Layer {
var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6))
var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
var flatten = Flatten<Float>()
var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)
@differentiable
func call(_ input: Tensor<Float>) -> Tensor<Float> {
return input.sequenced(through: conv, maxPool, flatten, dense)
}
}
畳み込み層（ conv ）のアクティベーションをバックプロパゲーション中に再計算する必要があります。ただし、 makeRecomputedInGradient(_:)を使用すると、特にsequenced(in:through:_:_:_:_:)を使用してレイヤーを順番に適用する場合に、結果のコードが煩雑に見える可能性があります。
input.sequenced(in: context, through: conv, maxPool, flatten, dense)
では、レイヤーをラップし、バックプロパゲーション中にそのアクティベーションを再計算する特別なレイヤータイプを定義してみませんか？やってみましょう。
まず、バイナリ関数をとるmakeRecomputedInGradient(_:)関数を定義します。
// Same as the previous `makeRecomputedInGradient(_:)`, except it's for binary functions.
func makeRecomputedInGradient<T: Differentiable, U: Differentiable, V: Differentiable>(
_ original: @escaping @differentiable (T, U) -> V
) -> @differentiable (T, U) -> V {
return differentiableFunction { x, y in
(value: original(x, y), pullback: { v in pullback(at: x, y, in: original)(v) })
}
}
次に、ジェネリックレイヤーActivationDiscarding<Wrapped>を定義します。
import TensorFlow
/// A layer wrapper that makes the underlying layer's activations be discarded during application
/// and recomputed during backpropagation.
struct ActivationDiscarding<Wrapped: Layer>: Layer {
/// The wrapped layer.
var wrapped: Wrapped
@differentiable
func callAsFunction(_ input: Wrapped.Input) -> Wrapped.Output {
let apply = makeRecomputedInGradient { (layer: Wrapped, input: Input) -> Wrapped.Output in
print(" Applying \(Wrapped.self) layer...")
return layer(input)
}
return apply(wrapped, input)
}
}
最後に、アクティベーションがアプリケーション中に破棄され、バックプロパゲーション中に再計算されることを除いて、同じレイヤーを返すすべてのレイヤーにメソッドを追加できます。
extension Layer {
func discardingActivations() -> ActivationDiscarding<Self> {
return ActivationDiscarding(wrapped: self)
}
}
モデルに戻ると、変更する必要があるのは、畳み込みレイヤーをアクティベーション再計算レイヤーにラップすることだけです。
var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()
今、それをモデルで使用するだけです！
struct Model: Layer {
var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()
var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))
var flatten = Flatten<Float>()
var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)
@differentiable
func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
return input.sequenced(through: conv, maxPool, flatten, dense)
}
}
トレーニングループを実行すると、畳み込みレイヤーのアクティブ化が2回計算されていることがわかります。1回はレイヤーの適用中、もう1回はバックプロパゲーション中です。
// Use random training data.
let x = Tensor<Float>(randomNormal: [10, 16, 16, 3])
let y = Tensor<Int32>(rangeFrom: 0, to: 10, stride: 1)
var model = Model()
let opt = SGD(for: model)
for i in 1...5 {
print("Starting training step \(i)")
print(" Running original computation...")
let (logits, backprop) = model.appliedForBackpropagation(to: x)
let (loss, dL_dŷ) = valueWithGradient(at: logits) { logits in
softmaxCrossEntropy(logits: logits, labels: y)
}
print(" Loss: \(loss)")
print(" Running backpropagation...")
let (dL_dθ, _) = backprop(dL_dŷ)
opt.update(&model, along: dL_dθ)
}
Starting training step 1
Running original computation...
Applying Conv2D<Float> layer...
Loss: 2.7602684
Running backpropagation...
Applying Conv2D<Float> layer...
Starting training step 2
Running original computation...
Applying Conv2D<Float> layer...
Loss: 2.5056229
Running backpropagation...
Applying Conv2D<Float> layer...
Starting training step 3
Running original computation...
Applying Conv2D<Float> layer...
Loss: 2.27868
Running backpropagation...
Applying Conv2D<Float> layer...
Starting training step 4
Running original computation...
Applying Conv2D<Float> layer...
Loss: 2.0752776
Running backpropagation...
Applying Conv2D<Float> layer...
Starting training step 5
Running original computation...
Applying Conv2D<Float> layer...
Loss: 1.8917936
Running backpropagation...
Applying Conv2D<Float> layer...
そのように、さまざまなドメインの一般的な微分可能プログラミングライブラリを定義するのは非常に簡単です。