Comments (6)
I believe this is an issue in Torchx itself and not in Axon. The Axon initializers are in the end all implemented through Defn, so if the backend options aren't being fetched and passed on properly, chances are this is related to Torchx.
An easy way to test this is to use EXLA in a machine with cuda support and set the default backend to {EXLA.Backend, client: :cuda}
.
from axon.
Thanks @polvalente
I have investigated this issue a bit more and it seems like in the initiation process a tensor is being created using Nx.BinaryBacked
and then this tensor get's passed to from_nx function in Torchx.Backend. Now based on torchx/backend.ex#L1808, Torchx will call Nx.backend_transfer
but it does not know about the default options anymore and will create cpu tensor instead of mps or any other device that the user has specified.
I can see two possible solutions:
- Investigate why the initializer is making a tensor in
Nx.BinaryBackend
and force it to useNx.default_backend
instead. - Create a new configuration parameter for Torchx which stores the default options for when the
from_nx
function wants to callbackend_transfer(tensor, {Torchx.Backend, default_options})
in torchx/backend.ex#L1808
from axon.
I just ran into the same issue - did you ever try out either of your two proposed solutions?
from axon.
Hi @benswift
I actually changed a line in Torchex Backend to forcefully create tensors with the mps device for my test and it worked as expected.
But I don't yet understand why tensors are created in BinaryBackend in the first place...
from axon.
Ah, ok, thanks for the update :)
from axon.
Yes, this is a torchx issue, so closing for now!
from axon.
Related Issues (20)
- Online deep learning in Axon HOT 2
- Compiled model will always result in `%Axon.None{}`
- Block with multiple inputs fails to compile HOT 2
- mixed precision policies cause bumblebee models to fail HOT 2
- Cannot return map from Axon.nx
- batch_norm states "mean" and "var" never updated HOT 1
- Optimize frozen parameter memory usage HOT 1
- Allow overriding layer metadata such as `:op` and `:op_name` for built-ins HOT 1
- XLA Slow operation alarms: "The operation took ..." HOT 1
- a
- Loss increasing with binary_cross_entropy above version 0.4.1 HOT 6
- Compile Error due to a type mismatch HOT 2
- Req usage in OTP 26 HOT 6
- Most `Axon.ModelState` functions should be `defn` or `deftransform`
- Merge model state updates into a single `Axon.ModelState` call
- Optionally return gradients / gradient state in train step
- Crash with mixed precision with new Nx
- Remove namespaces
- Can't pass multiple inputs in custom trainer functions HOT 3
- Ambiguous inputs should raise
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from axon.