Google/Trax – Understand and explore advanced deep learning

by Bella-Xiang • 6 years ago

359 50 6 years 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• 6 years ago

Trax is built on top of Jax, which is built on top of Autograd, which is a PyTorch like thing

sandGorgon2• 6 years 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• 6 years ago

Looking forward to a readme that is properly filled out. Some documentation as well. Looks promising.

codingslave2• 6 years 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• 6 years ago

I thought Theano was the first real deep learning implementation. What is the difference between that and tensorflow ?

codingslave2• 6 years ago

True, I was more so referencing first implementation to see wide scale adoption. But maybes there some holes in my narrative...

jhj• 6 years 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

m0zg6• 6 years 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• 6 years 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.

sseveran• 6 years ago

I agree with this, although I think it was a conscious and deliberate choice with TF 2.0. We have given up on TF for all future work which is sad since I really appreciate a number of the pieces that surround the core. I think they made a choice to emphasize the support of already developed models and make the experience great for novices will be a decision that they will come to regret. We found so many issues when we tried to port some of our existing models to TF 2.0. The sad part was that there were GitHub issues for all of them.

Personally I think Tensorflow has already lost and we just need to let it play out over the next few years. One interesting wrinkle is that since Trax, Jax and Flax utilize pieces of Tensorflow the TF team can probably claim good internal adoption numbers depending on how they count.

lowdose• 6 years ago

Every obscure solution in Tensorflow also has a change to break at an upgrade, I'm glad I moved to pytorch.

logicchains3• 6 years ago

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.

I'd agree if Google didn't have a history of building things with (arguably) unnecessarily complex APIs, like Angular1. I remember when Angular and React were new, seeing an Angular "cheatsheat" that was around 14 pages; the equivalent React cheatsheet was only 2 pages. Now, I do love the idea of Tensorflow 1, essentially a functional DSL to explicitly construct computation graphs, but Google's implementation of that idea was suboptimal: hard-to-follow error messages, not intuitive, multiple APIs to do the same thing (and continual API breakage as new APIs are introduced), difficult to debug. And even if the graph "compiled" correctly, it could still fail at execution time.

It's like they were building a programming language but lacked anyone with language design or PL theory background. Which makes sense given that anyone passionate about language design might prefer to work for somewhere closer to the cutting edge like Microsoft (C#, F#, Typescript, F*...), Facebook (Bucklescript, Hack), Apple (Swift) or Mozilla (Rust). Google does have its own languages, Dart and Go, but they're notable for ignoring and rejecting respectively cutting-edge PL theory (e.g. not disallowing null pointers, and in Go's case not even supporting parametric polymorphism). The day-to-day languages used at Google are also not particularly appealing to a PL enthusiast: Python, Java and non-modern C++.

Google software often also seems to care more about enforcing their idea of "best practices" on the user than about user experience. Tensorflow's C++ support is an example of this: it requires using Babel. As Babel doesn't easily support integration into an existing C++ project, this essentially means you have to change your whole project over to Babel just to use Tensorflow C++, which is a huge amount of effort to go to just to use a library. Especially when it's probably quicker to just rewrite the model in PyTorch, which provides a simple header file and static library for linking, the standard way of distributing C++ libraries. PyTorch also provides a nicer C++ API, because it's not confined to the ancient C++ standards that Google enforces, so can provide a modern API that's not full of macros (macros in modern C++ are considered bad practice; should only be used when there's absolutely no alternative).

To be fair, Google is really good at engineering language runtimes. Dart, Go and Tensorflow are all impressive works of engineering. They just seem to lack the organisational DNA for making them really nice to use, which maybe makes sense given the main source of their revenues is search/AdWords, the success of which is primarily driven by superior science/engineering. Compared to e.g. Facebook and Microsoft, that were/are in the business of making pretty things that users like to use (operating system, word processor, website). Or even comparing to Apple: people pay a huge premium for Apple phones over Android phones, in spite of their worse hardware, because of their more appealing design.

WnZ39p0Dgydaz14• 6 years ago

I would actually argue that the very first version of Tensorflow was great. I don't quite remember what version number that was, but it was well before 1.0. It was basically just a better version of Theano with a single way of doing things and very explicit graphs. I remember I ported all my Theano code without issues and loved it.

From then on it all went downhill. Lots of duplicate APIs were added and the code grew increasingly complex and opque because everyone at Google wanted in. Everyone wanted a slice of the pie to pad their resumes and internal performance reviews to be able to say "Look, I contributed this to TF! Promote me!" - Now it's a typical example of Conway's Law - a big mess that mirrors how messed up the internal incentive structure at Google is.

32gbsd• 6 years ago

Exactly what I call the programmer open-source API trap. The moment I saw how tensor was built out and published almost instantly I knew it was a funnel for capturing innovation and would eventually be replaced when they found a better mouse trap. I did not spend much time learning that open source wheel. Hopefully this new one will be smaller.

thedance• 6 years ago

It's a mistake to imagine that Google has some kind of centralized "programming languages" organization[1]. The people who designed Angular have probably never even met the people who designed Go, much less been influenced by them or shared influences. Same for TensorFlow. There may be "PL enthusiasts" at Google but the things they do at work are the same things everybody does at Google: moving data from one protobuf to another.

1: There is a "language platforms and tools" group (I used to work there) but its mission is stewardship of existing languages, not development of new ones.

DrScientist• 6 years ago

Somewhat like Map/Reduce was quickly replaced by Spark, which has no real competitors.

I wouldn't say that - what about Flink? (https://flink.apache.org/)

sandGorgon2• 6 years ago

Actually Dask on the python side.

Much higher performance if you were doing pyspark.

Because of the serialisation cost

BenGosub• 6 years ago

This is an example of Peter Thiel's "the last mover advantage" principle at work.

joaogui1• 6 years ago

Any idea about flax? And how do flax and Trax interact?

zelly• 6 years ago

That's a shame, I liked tensorflow much better than the imperative alternatives. I guess it comes with the territory. Scientific software has to be dumbed down on the CS side because the users (grad students) typically never seriously programmed.

sillysaurusx2• 6 years 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-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

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• 6 years ago

You can use tensorflow's tf.data, I think.

sillysaurusx2• 6 years 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.

zwaps2• 6 years 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?

ma2rten2• 6 years 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."

zwaps2• 6 years 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.

WnZ39p0Dgydaz14• 6 years ago

You're not alone. Even as an ML researcher at ex-FAANG I have no idea what this is. Is it a collection of well-documented models built on top of jax?

I could probably figure out what exactly this is if I spend an hour looking through the code, but it should be made clear in the README.

It's kind of funny, I think it's completely opaque what this actually is to 99% of HN users, but for some reason it's being upvoted because it has Google and Deep Learning in the name.

WnZ39p0Dgydaz14• 6 years ago

15min of looking through the code and my best guess right now is that it's a reimplementation of higher-level primitives, such as optimizers and layers, found in Tensorflow/Pytorch/etc, but based on a variable backend (you can pick jax or TF), together with a collection of models and training loops. I think the idea is that most of the code is simpler and more modular than what you would find in TF, which makes the models easier to read.

However, I don't yet understand what the use case is, or how it helps you to "learn" anything.

chillee5• 6 years ago

Trax is meant as a successor to Tensor2Tensor.

From my understanding of what Tensor2Tensor was (and thus what Trax is):

Basically, at it's core, it's a library of models, datasets, and utilities that are meant to be clear and work together nicely.

This is meant to be good for education, as you have an easy and unified way to run SOTA trained models on a wide variety of datasets.

This is also meant to be good for research, as much of research involves building on top of other people's models, and having a library of models with shared API (as opposed to the fragmented wilderness that is public research code) is meant to facilitate that.

I believe that it also ended up being used in industry as well, simply as a reliable source of pre-trained models.

PS: As Jax did not have an "official" neural network library until very recently, this also served as a neural network library for Jax.

shoyer• 6 years ago

JAX still doesn't have an "official" neural network library yet.

Source: I work closely with the JAX team at Google.

m0zg6• 6 years 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.

chillee5• 6 years 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).

m0zg6• 6 years 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• 6 years 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

Google/Trax – Understand and explore advanced deep learning

Related Stories