def testSimpleTraining(self):
def my_model(x, y):
x = layers.Conv2D(8, 3, padding='same', name="conv1", use_bias=False)(x)
x = layers.Conv2D(8, 3, padding='same', name="conv2", use_bias=False)(x)
x = layers.Conv2D(8, 3, padding='same', name="conv3", use_bias=False)(x)
x = math_ops.reduce_max(x, axis=[1, 2])
cross_entropy = nn.softmax_cross_entropy_with_logits_v2(
logits=x, labels=array_ops.stop_gradient(y))
loss = math_ops.reduce_mean(cross_entropy)
optim = so.ShardedOptimizer(gd.GradientDescentOptimizer(0.01))
train = optim.minimize(cross_entropy)
return [loss, train]
with ops.device("cpu"):
inp = array_ops.placeholder(np.float32, [1, 12, 12, 4], name="data")
lab = array_ops.placeholder(np.float32, [1, 8], name="labl")
with ops.device("/device:IPU:0"):
l, t = my_model(inp, lab)
autoshard.automatic_sharding(2, inp, l)
op_set = sharding.dependencies([l, t])
for o in op_set:
if o.device == '/device:IPU:0' and o.type not in allowed_op_types:
self.assertTrue(o.get_attr('_XlaSharding') is not None)
def testSimpleXlaCompileInference(self):
def my_model(inp):
output = inp * inp
return [output]
with ops.device("cpu"):
inp = array_ops.placeholder(np.float32, [], name="a")
with ops.device("/device:IPU:0"):
out = ipu_compiler.compile(my_model, inputs=[inp])
autoshard.automatic_sharding(2, inp, out[0])
op_list = ops.get_default_graph().get_operations()
for o in op_list:
if o.device == '/device:IPU:0' and o.type != 'NoOp':
self.assertTrue(o.get_attr('_XlaSharding') is not None)
def testMarkOpsWithAutoshardingContext(self):
with ops.device("/device:IPU:0"):
with autoshard.ipu_autoshard():
x = array_ops.placeholder(dtypes.float32, shape=[1, 32, 32, 4])
y = array_ops.placeholder(dtypes.float32, shape=[1, 8])
inp = x
with ops.name_scope('gradients'):
x = layers.Conv2D(
8, 3, padding='same', name="conv1", use_bias=False)(x)
x = layers.Conv2D(
8, 3, padding='same', name="conv2", use_bias=False)(x)
x = layers.Conv2D(
8, 3, padding='same', name="conv3", use_bias=False)(x)
x = math_ops.reduce_max(x, axis=[1, 2])
cross_entropy = nn.softmax_cross_entropy_with_logits_v2(
logits=x, labels=array_ops.stop_gradient(y))
loss = math_ops.reduce_mean(cross_entropy)
optim = so.ShardedOptimizer(gd.GradientDescentOptimizer(0.01))
optim.minimize(loss)
autoshard.automatic_sharding(2, inp, loss)
to_autoshard = ops.get_default_graph().get_collection(
sharding._IPU_AUTOSHARD)
fwd_ops = []
bwd_ops = []
all_ops = ops.get_default_graph().get_operations()
for o in all_ops:
if o in to_autoshard:
fwd_ops.append(o)
else:
bwd_ops.append(o)
self.assertTrue(len(fwd_ops) > 10)
self.assertTrue(len(bwd_ops) > 10)
self.assertEqual(len([o for o in fwd_ops if 'Conv2D' == o.type]), 3)
def my_model(loss, x, y):
with ops.device("/device:IPU:0"):
inp = x
lstm_cell = popnn_rnn.PopnnLSTM(256, dtype=dtypes.float32)
x, _ = lstm_cell(x, training=True)
cross_entropy = nn.softmax_cross_entropy_with_logits_v2(
logits=x, labels=array_ops.stop_gradient(y))
loss = math_ops.reduce_mean(cross_entropy)
optim = so.ShardedOptimizer(gd.GradientDescentOptimizer(0.01))
train = optim.minimize(cross_entropy)
autoshard.automatic_sharding(2, inp, loss)
return [loss, train]
def my_model(inp, lab):
x = inp
y = lab
x = layers.Conv2D(8, 3, padding='same', name="conv1", use_bias=False)(x)
x = layers.Conv2D(8, 3, padding='same', name="conv2", use_bias=False)(x)
x = layers.Conv2D(8, 3, padding='same', name="conv3", use_bias=False)(x)
x = math_ops.reduce_max(x, axis=[1, 2])
cross_entropy = nn.softmax_cross_entropy_with_logits_v2(
logits=x, labels=array_ops.stop_gradient(y))
loss = math_ops.reduce_mean(cross_entropy)
optim = so.ShardedOptimizer(gd.GradientDescentOptimizer(0.01))
train = optim.minimize(cross_entropy)
autoshard.automatic_sharding(2, inp, loss)
return [loss, train]
def my_model(lr, loss, x, y):
with ipu.ops.ipu_scope("/device:IPU:0"):
inp = x
x = layers.Conv2D(8,
3,
padding='same',
name="conv1",
use_bias=False)(x)
x = math_ops.reduce_max(x, axis=[1, 2])
cross_entropy = nn.softmax_cross_entropy_with_logits_v2(
logits=x, labels=array_ops.stop_gradient(y))
loss = math_ops.reduce_mean(cross_entropy)
optim = so.ShardedOptimizer(
gd.GradientDescentOptimizer(lr))
train = optim.minimize(cross_entropy)
autoshard.automatic_sharding(2, inp, loss)
return [lr, loss, train]