So it appears that when I run my code on Colab with GPU enabled it takes the same time as when I have just CPU. I’ve tried running my code inside the withDevice method like this
let (loss, grad) = withDevice(.gpu) {
return layer.valueWithGradient { layer -> Tensor<Float> in
.....
return loss
}
}
But I get the following error which seems to indicate it isn’t using the GPU as it can’t find it. I’ve also tried with withDefaultDevice but the code is still as slow as CPU
Fatal error: Device /job:localhost/replica:0/task:0/device:GPU:0 not found: file /tmp/tmp1habptzl/swift-install/package/.build/checkouts/swift-apis/Sources/TensorFlow/Core/Runtime.swift, line 1030
Current stack trace:
0 libswiftCore.so 0x00007fce777b88b0 swift_reportError + 50
1 libswiftCore.so 0x00007fce77827aa0 _swift_stdlib_reportFatalErrorInFile + 115
2 libswiftCore.so 0x00007fce7774face <unavailable> + 3738318
3 libswiftCore.so 0x00007fce7774fc47 <unavailable> + 3738695
4 libswiftCore.so 0x00007fce7751dc4d <unavailable> + 1436749
5 libswiftCore.so 0x00007fce77724a78 <unavailable> + 3562104
6 libswiftCore.so 0x00007fce7751d0a9 <unavailable> + 1433769
7 libjupyterInstalledPackages.so 0x00007fce74a2b4b0 _ExecutionContext.withDevice<A>(named:perform:) + 1187
8 libjupyterInstalledPackages.so 0x00007fce74a2afb0 _ExecutionContext.withDevice<A>(_:_:perform:) + 1110
9 libjupyterInstalledPackages.so 0x00007fce749dfcf0 withDevice<A>(_:_:perform:) + 169
11 repl_swift 0x0000000000400490 <unavailable> + 1168
13 libswiftCore.so 0x00007fce7751d0a9 <unavailable> + 1433769
14 libjupyterInstalledPackages.so 0x00007fce74a2b4b0 _ExecutionContext.withDevice<A>(named:perform:) + 1187
15 libjupyterInstalledPackages.so 0x00007fce74a2afb0 _ExecutionContext.withDevice<A>(_:_:perform:) + 1110
16 libjupyterInstalledPackages.so 0x00007fce749dfcf0 withDevice<A>(_:_:perform:) + 169
18 repl_swift 0x0000000000400490 <unavailable> + 1168
Current stack trace:
frame #2: 0x00007fce74a2b953 libjupyterInstalledPackages.so`_ExecutionContext.withDevice<R>(name="/job:localhost/replica:0/task:0/device:GPU:0", body=0x00007fce0d52d440 $__lldb_expr160`partial apply forwarder for reabstraction thunk helper from @callee_guaranteed () -> (@owned TensorFlow.Tensor<Swift.Float>, @owned __lldb_expr_83.Generator.AllDifferentiableVariables, @error @owned Swift.Error) to @escaping @callee_guaranteed () -> (@out (TensorFlow.Tensor<Swift.Float>, __lldb_expr_83.Generator.AllDifferentiableVariables), @error @owned Swift.Error) at <compiler-generated>, self=0x0000000003bd7960) at Runtime.swift:1030:13
frame #3: 0x00007fce74a2b406 libjupyterInstalledPackages.so`_ExecutionContext.withDevice<R>(kind=gpu, index=0, body=0x00007fce0d52d440 $__lldb_expr160`partial apply forwarder for reabstraction thunk helper from @callee_guaranteed () -> (@owned TensorFlow.Tensor<Swift.Float>, @owned __lldb_expr_83.Generator.AllDifferentiableVariables, @error @owned Swift.Error) to @escaping @callee_guaranteed () -> (@out (TensorFlow.Tensor<Swift.Float>, __lldb_expr_83.Generator.AllDifferentiableVariables), @error @owned Swift.Error) at <compiler-generated>, self=0x0000000003bd7960) at Runtime.swift:1024:20
frame #4: 0x00007fce749dfd99 libjupyterInstalledPackages.so`withDevice<R>(kind=gpu, index=0, body=0x00007fce0d52d440 $__lldb_expr160`partial apply forwarder for reabstraction thunk helper from @callee_guaranteed () -> (@owned TensorFlow.Tensor<Swift.Float>, @owned __lldb_expr_83.Generator.AllDifferentiableVariables, @error @owned Swift.Error) to @escaping @callee_guaranteed () -> (@out (TensorFlow.Tensor<Swift.Float>, __lldb_expr_83.Generator.AllDifferentiableVariables), @error @owned Swift.Error) at <compiler-generated>) at Execution.swift:37:41
frame #5: 0x00007fce0d523495 $__lldb_expr160`main at <Cell 25>:54:26
I’ve also tried the enableGPU function but it doesn’t seem to be valid on Colab. I have a feeling maybe it was added for .5 and Colab is running .4
Anyone have any idea what I’m doing wrong??
Thanks!
Here’s how I have the notebook set up.
