tape.py

with parameters() as params:
  fc1 = layers.Dense(hidden_dim, input_shape=(input_dim,))
  fc2 = layers.Dense(output_dim, input_shape=(hidden_dim,))

def forward(images, labels):
  x = fc1(images)
  x = layers.relu(x)
  x = fc2(x)
  logits = layers.relu(x)
  loss = losses.softmax_cross_entropy(logits, labels)
  return loss

forward_and_backward = grad(forward, params)

opt = tf.train.GradientDescentOptimizer(learning_rate=0.01)
for i in range(50):
  images, labels = image_batch()
  grads = forward_and_backward(images, labels)
  opt.apply_gradients(params, grads)

with parameters() as params:
  fc1 = layers.Dense(hidden_dim, input_shape=(input_dim,))
  fc2 = layers.Dense(output_dim, input_shape=(hidden_dim,))

def forward(images, labels):
  x = fc1(images)
  x = layers.relu(x)
  x = fc2(x)
  logits = layers.relu(x)
  return losses.softmax_cross_entropy(logits, labels)

opt = tf.train.GradientDescentOptimizer(learning_rate=0.01)
for i in range(50):
  images, labels = image_batch()
  tape.train(forward, params, opt, images, labels)

print(params)

One possible implementation of grad

def grad(forward, params):
  def forward_and_backward(*args, *kwargs):
    with tf.GradientTape() as grad_tape:
      loss = forward(*args, *kwargs)
    return grad_tape.gradient(loss, params)

  return forward_and_backward