def testVariableMapItems(self):
hidden = tf.ones(shape=(1, 16, 16, 3))
partitioner = tf.variable_axis_size_partitioner(4)
conv = snt.Conv2D(output_channels=3,
kernel_shape=3,
stride=1,
partitioners={"w": partitioner})
conv(hidden)
variable_map = snt.get_normalized_variable_map(conv)
items = snt.variable_map_items(variable_map)
items_str = sorted((key, var.op.name) for key, var in items)
self.assertEqual(items_str, [(u"b", u"conv_2d/b"),
("w", u"conv_2d/w/part_0"),
("w", u"conv_2d/w/part_1"),
("w", u"conv_2d/w/part_2")])
def testVariableMapItems(self):
hidden = tf.ones(shape=(1, 16, 16, 3))
partitioner = tf.variable_axis_size_partitioner(4)
conv = snt.Conv2D(output_channels=3,
kernel_shape=3,
stride=1,
partitioners={"w": partitioner})
conv(hidden)
variable_map = snt.get_normalized_variable_map(conv)
items = snt.variable_map_items(variable_map)
items_str = sorted((key, var.op.name) for key, var in items)
self.assertEqual(items_str, [(u"b", u"conv_2d/b"),
("w", u"conv_2d/w/part_0"),
("w", u"conv_2d/w/part_1"),
("w", u"conv_2d/w/part_2")])