8 ms·

Google/Trax – Understand and explore advanced deep learning

• deleted 7y ago
  [deleted]
• unityByFreedom 7y ago
  Is this like a layer on top of TensorFlow to make it easier to get started? Is it meant to compete with PyTorch in that respect? I wish the title and description were more clear. They make it sound like a course but it is a library/command-line tool.
  • chewxy 7y ago
    Trax is built on top of Jax, which is built on top of Autograd, which is a PyTorch like thing
  • sandGorgon 7y ago
    So Tensorflow 2 is built on top of a Keras API. Its supposed to be better UX, but Pytorch really is far superior UX.
• JPKab 7y ago
  Looking forward to a readme that is properly filled out. Some documentation as well. Looks promising.
• codingslave 7y ago
  Word is that internally at Google, among a few teams, and then also externally, Trax/Jax are putting up real competition to Tensorflow. Some teams have moved off of tensorflow entirely. Combined with the better research capabilities of PyTorch, the future of tensorflow is not bright. Given that, Tensorflow still provides the highest performance with regards to production usage, and has tons of legacy code strewn throughout the web. I would argue that this is not the fault of Tensorflow, but rather the hazard of being the first implementation in an extremely complex space. Seems like usually there needs to be some sacrificial lamb in software domains. Somewhat like Map/Reduce was quickly replaced by Spark, which has no real competitors.
  • zitterbewegung 7y ago
    I thought Theano was the first real deep learning implementation. What is the difference between that and tensorflow ?
    • jhj 7y ago
      The first would be SN/Lush from circa 1987, of which Torch (from NEC Research), and then PyTorch are direct descendants in a lot of ways. https://leon.bottou.org/projects/lush https://leon.bottou.org/projects/lush
    • codingslave 7y ago
      True, I was more so referencing first implementation to see wide scale adoption. But maybes there some holes in my narrative...
      • danieldk 7y ago
        Actually, Torch (2002) is older than Theano (2007). Moreover, both were already used widely in the research community. Tensorflow was actually late in the grand scheme of things, but replaced Theano and to some extend Torch (which used Lua for its public API). I think it's fairer to say that Tensorflow is the first implementation that had wide adoption outside academia.
  • m0zg 7y ago
    >> highest performance with regards to production usage Not really, no. I've been using TensorRT for that quite successfully. If you can work around its limitations, I don't think anything can compete, at least not on modern NVIDIA GPUs. And yes, I'm aware TF can use TensorRT as well. But why drag in all of TF if you just want inference?
  • maiybe 7y ago
    Overall, I've seen similar movement away from Tensorflow in my social circle of research scientists/engineers. One area I'd push back on is that "this is not the fault of Tensorflow." An area of weakness for Tensorflow is that it solves a number of DL problems with a specialized API call. That's not an asset, that's a liability. LSTMs were always a pain point. So much so that for Tensorflow projects, I gave up and insisted on traditional feedforward approaches like CNNs + MLPs or ResNets when LSTMs would be viable. Mostly identical performance with decent speed boosts from avoiding recurrence, and the simpler code reduced maintenance by non-ML engineers. As soon as you branch out of standard DL bread and butter models, you spend frustratingly long periods of time tracking down obscure solutions in a part of the API space that had its own hard-to-follow logic. Every time I'd point out that it's hard to do something either in forums or HN directly, I'd get a response that its easy to do with [insert-random-api] function call. In the end, it's my opinion that Tensorflow will lose out to JAX and Pytorch, by no fault other than its own complicated construction.
• sillysaurusx 7y ago
  I was recently surprised to discover that Jax can't use a TPU's CPU, and that there are no plans to add this to Jax. https://github.com/google/jax/issues/2108#issuecomment-581541862 https://github.com/google/jax/issues/2108#issuecomment-58154... A TPU's CPU is the only reason that TPUs are able to get such high performance on MLPerf benchmarks like imagenet resnet training. https://mlperf.org/training-results-0-6 https://mlperf.org/training-results-0-6 They do infeed processing (image transforms, etc) on the TPU's CPU. Then the results are fed to each TPU core. Without this capability, I don't know how you'd feed the TPUs with data in a timely fashion. It seems like your input will be starved. Hopefully they'll bring jax to parity with tensorflow in this regard soon. Otherwise, given that jax is a serious tensorflow competitor, I'm not sure how the future of TPUs will play out. (If it sounds like this is just a minor feature, consider how it would sound to say "We're selling this car, and it can go fast, but it has no seats." Kind of a crucial feature of a car.) Still, I think this is just a passing issue. There's no way that Google is going to let their TPU fleet languish. Not when they bring in >$1M/yr per TPU pod commitment.
  • yablak 7y ago
    You can use tensorflow's tf.data, I think.
    • sillysaurusx 7y ago
      Sort of. The way you do it is to scope operations to tf.device(None), which selects the TPU's CPU. (I think it's equivalent to using the first device in sess.list_devices(), which is the CPU.) You can scope operations in tf.data, sure, but you can also execute arbitrary training operations. I use this technique to train GPT-2 117M with a 25k context window, which requires about 300GB memory. Only a TPU's CPU has so much. That's why it was surprising to hear Jax can't do that. It's one of the best features of a TPU.
• zwaps 7y ago
  Is it just me or is there zero explanation to what this actually is? It somehow "helps" me understand deep learning but its tutorial / doc is one python notebook with three cells where some nondescript unknown API is called to train a transformer. Huh?
  • ma2rten 7y ago
    From the README: "Trax code is structured in a way that allows you to understand deep learning from scratch. We start with basic maths and go through layers, models, supervised and reinforcement learning."
    • zwaps 7y ago
      But what is it? It says it is an API or module? Is it a DL framework? Is it an extension of TF/Jax with new models? Is it a set of tutorials? Are we to work through the code? Is a part of the doc for Jax? Is it a set of implementations of Jax code into concrete models? All these things come to mind when going through this Github. By now I think to have figured out that it is a module that implements some DL models using Jax. So it is like an extension to Jax. It "helps" you to understand how to build the models in Jax if you go through the code. You could also load these abstractions as a Python module and use them in code, but I would doubt that this would "help you understand deep learning". So the real intention is to go through the code and see how Jax implements, for example, a transformer. Alternatively, this could be an early stage of some PyTorch-type framework over Jax and TF. So essentially, another Keras.
      • kavalg 7y ago
        Thank you
      • ma2rten 7y ago
        It's the successor of tensor2tensor, so it's a high level library that implements different the algorithms. Unlike other libraries the code is very readable.
  • WnZ39p0Dgydaz1 7y ago
• m0zg 7y ago
  Not sure why one would bother with this. This is a less mature version of PyTorch. And I know there's XLA and stuff, but I've yet to see any major benefit from that for research in particular. A ton of time in DL frameworks is spent in the kernels (which in most practical cases means CUDA/cuDNN) which are hand-optimized far better than anything we'll ever get out of any optimizer.
  • chillee 7y ago
    If you're talking about Jax, there's a couple different reasons to bother for research 1. Full numpy compatibility. 2. More efficient higher order gradients (because of forward mode auto diff). Naively it's asymptotic improvement, but I believe Pytorch uses some autodiff tricks to perform higher order gradients with backwards mode, at the cost of a decently high constant factor. 3. Some cool transformations like vmap. 4. Full code gen, which is neat especially for scientific computing purposes. 5. A neat API for higher order gradients. 2. and 5. are the most appealing for DL research, 1., 3. and 4. are appealing for those in the stats/scientific computing communities. PyTorch is working on all of these, to various degrees of effort, but Jax currently has an advantage in these points (and may have a fundamental advantage in design in some).
    • m0zg 7y ago
      From a practitioner: 1. Meh. PyTorch is close enough to not worry about it, and is better in some places. 2. Meh. All the methods people use in practice for deep learning in particular do not use higher order gradients. Most higher order methods are prohibitively memory expensive, and memory is at a premium in acceleration hardware (and so is the bus bandwidth - so you can't "swap to RAM"). I do agree that higher order gradients are the next frontier in optimization though - current optimizer research seems to have stalled, so people focus on training with huge batches and stuff like that. Most SOTA models in my field are trained with SGD+momentum - super primitive stuff. I don't see how Jax would solve the memory problem though. You still have to store those Hessians somewhere, at least partially. 3. Do agree, that's cool if it actually parallelizes nontrivial stuff which e.g. tf.vectorized_map barfs on. Although in a lot of cases you can "vectorize" by concatenating input tensors into a higher dimensional tensor. 4. Meh. Not sure why I'd want that if I already have tracing and JIT. 5. This is #2 With PyTorch though, you get close enough to Numpy to feel at home in both, and there's so much code written for it already that you can usually find a good starting point for your research pretty easily on github and then build on top of that. If you need to deploy, there's also tracing and jit, which lets you load and serve models with libtorch. I see what you're saying regarding "advantages", I'm just pointing out that PyTorch might be "good enough" for most people. If I were on that team, I'd focus on providing comfortable transition from TF 2.x which is a dumpster fire (with the exception of TensorFlow Lite which is excellent). That, IMO, would be the only way for this project to achieve mainstream success unless PyTorch disintegrates over time.
• nestorD 7y ago
  Note that, in this space, there is also Flax[0] which is also built on top of Jax bringing more deep-learning specific primitives (while not trying to be tensorflow compatible unlike Trax if I understand correctly). [0]: https://github.com/google-research/flax/tree/prerelease https://github.com/google-research/flax/tree/prerelease