def _CreateChildrenVariables(self):
with tf.variable_scope(self.params.name):
# Prepends p.num_experts to the tensor shape of every variable created
# by p.body.
with py_utils.VariableShapePrefixContext(self.params.num_experts):
self.body.InstantiateVariables()
super()._CreateChildrenVariables()
def __init__(self, params):
super(RepeatLayer, self).__init__(params)
p = self.params
assert p.name
assert p.repeat > 0
with tf.variable_scope(p.name):
with py_utils.VariableShapePrefixContext(p.repeat):
self.CreateChild('body', p.body)
def __init__(self, params):
super().__init__(params)
p = self.params
assert p.name
assert p.repeat > 0
if p.per_layer_vars:
for i in range(p.repeat):
self.CreateChild('body_iter_%05d' % i, p.body)
else:
with py_utils.VariableShapePrefixContext(p.repeat):
self.CreateChild('body', p.body)
def _CreateChildrenVariables(self):
p = self.params
with tf.variable_scope(self.params.name):
if p.per_layer_vars:
for i in range(p.repeat):
with tf.variable_scope('iter_%05d' % i):
self.children['body_iter_%05d' % i].InstantiateVariables()
else:
with py_utils.VariableShapePrefixContext(self.params.repeat):
self.body.InstantiateVariables()
super()._CreateChildrenVariables()
def __init__(self, params):
super().__init__(params)
p = self.params
assert p.name
assert p.num_experts
assert p.cond_dim
with tf.variable_scope(p.name):
# Prepends p.num_experts to the tensor shape of every variable created
# by p.body.
with py_utils.VariableShapePrefixContext(p.num_experts):
self.CreateChild('body', p.body)
def __init__(self, params):
super(SoftCondLayer, self).__init__(params)
p = self.params
assert p.name
assert p.num_experts
assert p.input_dim
with tf.variable_scope(p.name):
# Create Variables for task weight mapping.
w_p = py_utils.WeightParams(
shape=[p.input_dim, p.num_experts],
init=p.params_init, # TODO(huangyp): try zero init instead.
dtype=p.dtype,
collections=[self.__class__.__name__ + '_vars'])
self.CreateVariable('w', w_p)
# Prepends p.num_experts to the tensor shape of every variable created
# by p.body.
with py_utils.VariableShapePrefixContext(p.num_experts):
self.CreateChild('body', p.body)
def _CreateChildrenVariables(self):
with tf.variable_scope(self.params.name):
with py_utils.VariableShapePrefixContext(self.params.repeat):
self.body.InstantiateVariables()
super()._CreateChildrenVariables()
def _CreateChildrenVariables(self):
with tf.variable_scope(self.params.name):
with py_utils.VariableShapePrefixContext(self.params.num_experts):
self.body.CreateVariables()
super()._CreateChildrenVariables()
