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]
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(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]