Skip to content

Instantly share code, notes, and snippets.

@rxwei
Created May 20, 2019 00:57
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save rxwei/e1488cac5379ba2bc3aff7490e18158f to your computer and use it in GitHub Desktop.
Save rxwei/e1488cac5379ba2bc3aff7490e18158f to your computer and use it in GitHub Desktop.
Broadcast operator gradient performance test
import TensorFlow
import Dispatch
func foo(x: Tensor<Float>, y: Tensor<Float>) -> Tensor<Float> {
return (x + x + x * y * y * y).sum()
}
func time(_ body: () -> Void) {
let divisor: Float = 1_000_000_000
let start = Float(DispatchTime.now().uptimeNanoseconds) / divisor
print("Start: \(start)s")
body()
let end = Float(DispatchTime.now().uptimeNanoseconds) / divisor
print("End: \(end)s")
let elapsed = end - start
print("Elapsed: \(elapsed)s")
}
let x = Tensor<Float>(zeros: [10, 1, 10, 1, 1])
let y = Tensor<Float>(zeros: [1, 10, 1, 10, 1])
time {
for i in 0..<10000 {
_ = gradient(at: x, y, in: foo)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment