yaroslavvb2

wget https://raw.githubusercontent.com/yaroslavvb/stuff/master/matmul_benchmark_seq.py
export TF_CPP_MIN_LOG_LEVEL=1
export CUDA_VISIBLE_DEVICES=0
python matmul_benchmark_seq.py --dtype=float16
python matmul_benchmark_seq.py --dtype=float32

yaroslavvb2

# expensive way to compute factorial of n
def factorial(n):
  def f(x):
    return tf.pow(x, n)
  for i in range(n):
    f = tfe.gradients_function(f)
  return f(1.)

yaroslavvb2

@tfe.custom_gradient
def noisy_square(x):
   def grad(b):
       true_grad = 2*b*x
       return true_grad+tf.random_uniform(())
   return (x*x), grad
grad = tfe.gradients_function(noisy_square)
x = 2.
points = []
for i in range(20):

yaroslavvb2

@graph_callable.graph_callable([])
def resnet_loss():
 “””Resnet loss from random input”””
 network = resnet_model.cifar10_resnet_v2_generator(RESNET_SIZE, NUM_CLASSES)
 inputs = tf.reshape(images, [BATCH_SIZE, HEIGHT, WIDTH, DEPTH])
 logits = network(inputs,True)
 cross_entropy = tf.losses.softmax_cross_entropy(logits=logits,
 onehot_labels=labels)
 return cross_entropy
loss_and_grads_fn = tfe.implicit_value_and_gradients(resnet_loss)

yaroslavvb2

import tensorflow as tf
from tensorflow.contrib.eager.python import tfe
tfe.enable_eager_execution()
context = tf.device('/gpu:0')
context.__enter__()

# download resnet_model
import sys, os, urllib.request
resnet_model_url="https://raw.githubusercontent.com/tensorflow/models/master/official/resnet/resnet_model.py"
response = urllib.request.urlopen(resnet_model_url)

yaroslavvb2

a = a.gpu() # copies tensor to default GPU (GPU0)
a = a.gpu(0) # copies tensor to GPU0
a = a.gpu(1) # copies tensor to GPU1
a = a.cpu() # copies tensor back to CPU

yaroslavvb2

pip install tf-nightly-gpu
python
from tensorflow.contrib.eager.python import tfe
tfe.enable_eager_execution()
a = tf.random_uniform((10,))
b = tf.random_uniform((10,))
for i in range(100):
  a = a*a
  if a[0]>b[0]:
  break

yaroslavvb2

mode = 'capture'
loss.backward()

# invert
for i in range(n):
  As_inv.append(regularized_inverse(As[i] @ As[i].t()))
  Bs_inv.append(regularized_inverse(Bs[i] @ Bs[i].t()))

mode = 'kfac'
loss.backward()

yaroslavvb2

def backward(ctx, grad_output):
  matrix1, matrix2 = ctx.saved_variables
  grad_matrix1 = grad_matrix2 = None

  if mode == 'capture':
    Bs.insert(0, grad_output.data)
    As.insert(0, matrix2.data)
  elif mode == 'kfac':
    B = grad_output.data
    A = matrix2.data