MARTIN GORNER: Hi, everyone, and thank you for being here

at 8:30 in the morning, and welcome to this session

about TPUs and TPU pods.

So those are custom made accelerators

that Google has designed to accelerate machine learning

workloads.

And before I tell you everything about them, me and Kaz,

I would like to do something.

Of course, this is live, so you want to see a live demo.

And I would like to train with you here, onstage, using

a TPU pod, one of those big models

that used to take days to train.

And we'll see if we can finish the training

within this session.

So let me start the training.

I will come back to explaining exactly what I'm doing here.

I'm just starting it.

Run all cells.

Seems to be running.

OK, I'm just checking.

I'm running this on a 128 core TPU pod.

So one of the things you see in the logs here

is that I have all my TPUs appearing.

0, 1, 2, 6, and all the way down to 128.

All right, so this is running.

I'm happy with it.

Let's hear more about TPUs.

So first of all, what is this piece of silicon?

And this is the demo that I've just launched.

It's an object detection demo that

is training on a wildlife data set of 300,000 images.

Why wildlife?

Because I can show you cute pandas.

And I can show you cute electronics as well.

So this is a TPU v2.

And we have a second version now, a TPU v3.

Those are fairly large boards.

It's large like this, roughly.

And as you can see, they have four chips on them.

Each chip is dual core, so each of these boards

has 8 TPU cores on them.

And each core has two units.

That is a vector processing unit.

That is a fairly standard data-oriented processor,

general purpose processor.

What makes this special for machine learning

is the matrix multiply unit.

TPUs have a built-in hardware-based matrix

multiplier that can multiply to 128 by 128 matrices in one go.

So what is special about this architecture?

There are two tricks that we used

to make it fast and efficient.

The first one is, I would say, semi-standard.

It's reduced precision.

When you train neural networks, reducing the precision

from 32-bit floating points to 16-bit

is something that people quite frequently do,

because neural networks are quite resistant to the loss

of precision.

Actually, it even happens sometimes

that the noise that is introduced by reduced precision

acts as a kind of regularizer and helps with convergence.

So sometimes you're even lucky when you reduce precision.

But then, as you see on this chart, float16 and float32,

the floating point formats, they don't have the same number

of exponent bits, which means that they

don't cover the same range.

So when you take a model and downgrade all your float32s

into float16s, you might get into underflow or overflow

problems.

And if it is your model, it's usually not

so hard to go in and fix.

But if you're using code from GitHub

and you don't know where to fix stuff,

this might be very problematic.

So that's why on TPUs, we chose a different--

actually we designed a different floating point

format called bfloat16.

And as you can see, it's actually

exactly the same as float32 with just the fractional bits

cut off.

So the point is it has exactly the same number

of exponent bits, exactly the same range.

And therefore, usually, it's a drop-in replacement

for float32 and reduced precision.

So typically for you, there is nothing

to do on your model to benefit from the speed of reduced

precision.

The TPU will do this automatically, on ship,

in hardware.

And the second trick is architectural.

It's the design of this matrix multiply unit.

So that you understand how this works,

try to picture, in your head, how to perform a matrix

multiplication.

And one result, one point of the resulting

matrix, try to remember calculus from school, is a dot product.

A dot product of one line of one matrix and one column

of the second matrix.

Now what is a dot product?

A dot product is a series of multiply-accumulate operations,

which means that the only operation you

need to perform a matrix multiplication

is multiply and accumulate.

And multiply-accumulate in 16 bits,

because we're using bfloat16 reduced precision.

That is a tiny, tiny piece of silicon.

A 16-bit multiply-accumulator is a tiny piece of silicon.

And if you wire them together as an array, as you see here.

So this in real life would be a 128 by 128 array.

It's called a systolic array.

Systolic in Greek means flow.

Because you will flow the data through it.

So the way it works is that you load one matrix into the array,

and then you flow the second matrix through the array.

And you'll have to believe me, or maybe

spend a little bit more time with the animation,

by the time the gray dots have finished

flowing through those multiply-accumulators,

out of the right side come all the dot products that

make the resulting matrix.

So it's a one-shot operation.

There are no intermediate values to store anywhere, in memory,

in registers.

All the intermediate values flow on the wires

from one compute unit to the second compute units.

It's very efficient.

And what is more, it's only made of

those tiny 16-bit multiply-accumulators,

which means that we can cram a lot of those into one chip.

128 by 128 is 16,000 multiply-accumulators.

And that's how much you get in one TPU core, twice that

in two TPU cores.

So this is what makes it dense.

Density means power efficiency.

And power efficiency in the data center means cost.

And of course, you want to know how cheap

or how fast these things are.

Some people might remember from last year

I did a talk about what I built, this planespotting model,

so I'm using this as a benchmark today.

And on Google Cloud's AI platform,

it's very easy to get different configurations,

so I can test how fast this trains.

My baseline is-- on a fast GPU, this model

trains in half in four and a half hours.

But I can also get 5 machines with powerful GPUs

in a cluster.

And on those five machines, five GPUs,

this model will train in one hour.

And I've chosen this number because one hour is exactly

the time it takes for this model to train on one TPU v2.

So the rule of thumb I want you to remember

is that roughly 1 TPU v2, with its 4 chips,

is roughly as fast as five powerful GPUs.

That's in terms of speed.

But as you can see, it's almost three times cheaper.

And that's the point of optimizing the architecture

specifically for neural network workloads.

You might want to know how this works in software as well.

So when you're using TensorFlow, or Keras in TensorFlow,

your Python code TensorFlow Python code

generates a computational graph.

That is how TensorFlow works.

So your entire neural network is represented as a graph.

Now, this graph is what is sent to the TPU.

Your TPU does not execute Python code.

This graph is processed through XLA, the Accelerated Linear

Algebra compiler, and that is how

it becomes TPU microcode to be executed on the TPU.

And one nice side-effect of this architecture

is that if, in your TensorFlow code,

you load your data through the standard tf.data.Dataset

API, as you should, and as is required with TPUs, then

even the data loading part, or imagery resizing, or whatever

is in your data pipeline, ends up in the graph,

ends up executed on the TPU.

And the TPU will be pulling data from Google Cloud Storage

directly during training.

So that is very efficient.

How do you actually write this with code?

So let me show you in Keras.

And one caveat, this is Keras in TensorFlow 1.14,

which should be out in these next days.

The API is slightly different in TensorFlow 1.13

today, but I'd rather show you the one that will be--

the new one, as of tomorrow or next week.

So it's only a couple of lines of code.

There is the first line, TPUClusterResolver.

You can call it without parameters on most platforms,

and that finds the connected TPU.

The TPU is a remotely-connected accelerator.

This finds it.

You initialize the TPU and then you use the new distribution

API in TensorFlow to define a TPU strategy based on this TPU.

And then you say with strategy.scope,

and everything that follows is perfectly normal Keras code.

Then you define your model, you compile it,

you do model.fit, model.evaluate, model.predict,

anything you're used to doing in Keras.

So in Keras, it's literally these four lines of code

to add--

to work on a TPU.

And I would like to point out that these four lines of code

also transform your model into a distributed model.

Remember a TPU, even a single GPU,

is a board with eight cores.

So from the get go it's distributed computing.

And these four lines of code put in place

all the machinery of distributed computing for you.

One parameter to notice.

You see in the TPU strategy, there

is the steps_per_run equals 100.

So that's an optimization.

This tells the TPU, please run 100 batches worth of training

and don't report back until you're finished.

Because it's a network attached accelerator,

you don't want the TPU to be reporting back

after each batch for performance reasons.

So this is the software.

If you don't want to write your own code,

I encourage you to do so.

But if you don't, we have a whole library

of TPU optimized models.

So you will find them on the TensorFlow/tpu GitHub

repository.

And there is everything in the image--

in the vision space, in the machine translation,

and language, and NLP space, in speech recognition.

Even you can play with GaN models.

The one that we are demoing on stage,

remember we are training the model right now, is RetinaNet.

So this one is an object detection model.

And I like this model, so let me say a few words

about how this works.

In object detection, you put an image, and what you get

is not just the label--

this is a dog, this is a panda--

but you actually get boxes around where those objects are.

In object detection models, you have two kinds.

There are one shot detectors that

are usually fast but kind of inaccurate,

and then two-stage detectors that are much more

accurate but much slower.

And I like RetinaNet because they actually

found a trick to make this both the fastest

and the most accurate model that you can

find in object detection today.