Help us understand the problem. What is going on with this article?

Swift for TensorFlowでVAE

オフィシャルのswift-modelsリポジトリにautoencoderのサンプルが掲載されています。
https://github.com/tensorflow/swift-models/blob/2fa11ba1d28ef09454af9da77e22b585cf3e5b7b/Autoencoder/main.swift

ですがVAEのサンプルはなかったので上コードを改造して実装してみました。

以下S4TF及びVAEの基本的なことは知っているものとして書きます。
またコードの一部はkerasの実装を参考にしました。
https://keras.io/examples/variational_autoencoder/

完成したコードはこちらです。v0.4.0-rc2にて動作確認しています。
https://gist.github.com/t-ae/3cd33e4f0535b98c2df9bfeef49645e5

モデル定義

struct Encoder: Layer {
    typealias Input = Tensor<Float>
    typealias Output = Encoded

    var encoder1 = Dense<Float>(inputSize: imageHeight * imageWidth, outputSize: 128,
            activation: relu)
    var encoder2 = Dense<Float>(inputSize: 128, outputSize: 64, activation: relu)
    var encoder3 = Dense<Float>(inputSize: 64, outputSize: 12, activation: relu)

    var encoderMean = Dense<Float>(inputSize: 12, outputSize: 4, activation: identity)
    var encoderLogVar = Dense<Float>(inputSize: 12, outputSize: 4, activation: identity)

    @differentiable
    func callAsFunction(_ input: Input) -> Output {
        let intermediate = input.sequenced(through: encoder1, encoder2, encoder3)

        let mean = encoderMean(intermediate)
        let logVar = encoderLogVar(intermediate)

        return Encoded(mean: mean, logVar: logVar)
    }
}

struct Encoded: Differentiable {
    var mean: Tensor<Float>
    var logVar: Tensor<Float>
}

struct Decoder: Layer {
    typealias Input = Tensor<Float>
    typealias Output = Tensor<Float>

    var decoder1 = Dense<Float>(inputSize: 4, outputSize: 12, activation: relu)
    var decoder2 = Dense<Float>(inputSize: 12, outputSize: 64, activation: relu)
    var decoder3 = Dense<Float>(inputSize: 64, outputSize: 128, activation: relu)
    var decoder4 = Dense<Float>(inputSize: 128, outputSize: imageHeight * imageWidth,
        activation: tanh)

    @differentiable
    func callAsFunction(_ input: Input) -> Output {
        return input.sequenced(through: decoder1, decoder2, decoder3, decoder4)
    }
}

struct VAE: Layer {
    typealias Input = Tensor<Float>
    typealias Output = VAEResult

    var encoder = Encoder()
    var decoder = Decoder()

    @differentiable
    func callAsFunction(_ input: Input) -> Output {
        let encoded = encoder(input)

        let mean = encoded.mean
        let logVar = encoded.logVar

        let gaussian = Tensor<Float>(randomNormal: mean.shape)

        let std = exp(logVar/2)

        let images = decoder(gaussian * std + mean)

        return VAEResult(image: images, mean: mean, logVar: logVar)
    }
}

struct VAEResult: Differentiable {
    var image: Tensor<Float>
    var mean: Tensor<Float>
    var logVar: Tensor<Float>
}

もともとあったAutoencoderの代わりにVAEを定義しています。
Autoencoderではエンコーダとデコーダが結合した形でしたが、ここでは分割し、それをVAEにて統合しています。これはPythonのフレームワークを使うときにも言えますが、分割しておくとエンコーダ/デコーダ単体で使いたいときに便利です。

VAEはサンプリングと損失の計算にエンコーダーから得られるmean, logVarを使うため、Encoderはこの2つを出力する必要があります。named tupleで出したいところですが@differentiableな関数の入出力はDifferentiableである必要があるため、ここではEncodedを定義しています。VAEResultについても同様です。

画像が[0, 1]範囲なのにDecoderの出力がtanhになっていますが、これは元のコードでもそうなっており、v0.4.0-rc2の時点ではsigmoidの勾配の実装にバグがあって正しく学習できないのでそのままにしているだけです。
https://github.com/tensorflow/swift-models/issues/179

損失関数

@differentiable
func loss(result: VAEResult, original: Tensor<Float>) -> Tensor<Float> {
    let reconstrcutionLoss = (result.image - original).squared().sum(alongAxes: 1)

    let klLoss = (1 + result.logVar - result.mean.squared() - exp(result.logVar))
        .sum(alongAxes: 1) * -0.5

    return (reconstrcutionLoss + klLoss).mean()
}

損失関数の実装はVAEResultと元画像のTensor<Float>を取ります。
var klLoss = (1 + result.logVar - result.mean.squared() - exp(result.logVar)); klLoss = klLoss.sum(alongAxes: 1); klLoss *= -0.5のようにvarに操作を加えていく書き方は@differentiableではできないようなので一文で書いています。毎回新しい変数を定義していけば分割はできますが……
(追記:どうも*=が問題なようで、klLoss = klLoss * -0.5を使うとコンパイルできました。)

学習部分

let 𝛁model = vae.gradient { vae -> Tensor<Float> in
    let result = vae(x)
    return loss(result: result, original: x)
}
optimizer.update(&vae.allDifferentiableVariables, along: 𝛁model)

損失関数は定義できているのであとは置きかえるだけです。

結果

50エポックやって最後の方の出力はこんなかんじになりました。
左が入力で右が出力です。

VAE出力

Swift for TensorFlowについて

まだ触りはじめたばかりですが、Pythonのフレームワークを使う場合と比べ圧倒的に書きやすいと感じます。
PythonのほうではPyCharmで開発していますが、補完で候補が出ないことがよくあるため、ドキュメントを検索しながらコードを書くということが多いです。一方S4TFではだいたいこんな名前だろ〜というのを打ったらちゃんと候補に出てくれるので楽です(Xcode/Swiftでの開発に慣れているというのも大きいかもしれませんが)。
さらに静的型付けのおかげで引数に何を渡せばいいかも明確なので、実装して走らせて初めて間違いに気付くというようなこともないです。

S4TFの正式リリースがいつになるかはわかりませんが、現状最も期待できる候補の一つなので、今後とも注視していきたいです。


後日GANも書いてこちらは本家に入れてもらいました。
https://github.com/tensorflow/swift-models/tree/master/GAN

Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
No comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
ユーザーは見つかりませんでした