experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #40619 from jpodivin:docfix37293

Browse Source 
PiperOrigin-RevId: 320625911
Change-Id: I02ad15e3b1c1ac97918db4b059939f8cc416da00
This commit is contained in:
TensorFlower Gardener 2020-07-10 10:29:48 -07:00
commit 7e279d6b0f
1 changed files with 31 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
tensorflow/python/module/module.py
View File

@ -39,19 +39,19 @@ class Module(tracking.AutoTrackable):
functions which apply to user input. For example a dense layer in a neural
network might be implemented as a `tf.Module`:
>>> class Dense(tf.Module):
... def __init__(self, in_features, out_features, name=None):
... super(Dense, self).__init__(name=name)
... self.w = tf.Variable(
... tf.random.normal([in_features, out_features]), name='w')
... self.b = tf.Variable(tf.zeros([out_features]), name='b')
... def __call__(self, x):
... y = tf.matmul(x, self.w) + self.b
... return tf.nn.relu(y)
>>> class Dense(tf.Module):
... def __init__(self, input_dim, output_size, name=None):
... super(Dense, self).__init__(name=name)
... self.w = tf.Variable(
... tf.random.normal([input_dim, output_size]), name='w')
... self.b = tf.Variable(tf.zeros([output_size]), name='b')
... def __call__(self, x):
... y = tf.matmul(x, self.w) + self.b
... return tf.nn.relu(y)
You can use the Dense layer as you would expect:
>>> d = Dense(in_features=3, out_features=2)
>>> d = Dense(input_dim=3, output_size=2)
>>> d(tf.ones([1, 3]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=..., dtype=float32)>
@ -77,22 +77,28 @@ class Module(tracking.AutoTrackable):
`with self.name_scope:` or you can annotate methods (apart from `__init__`)
with `@tf.Module.with_name_scope`.
```python
class MLP(tf.Module):
def __init__(self, input_size, sizes, name=None):
super(MLP, self).__init__(name=name)
self.layers = []
with self.name_scope:
for size in sizes:
self.layers.append(Dense(input_size=input_size, output_size=size))
input_size = size
>>> class MLP(tf.Module):
... def __init__(self, input_size, sizes, name=None):
... super(MLP, self).__init__(name=name)
... self.layers = []
... with self.name_scope:
... for size in sizes:
... self.layers.append(Dense(input_dim=input_size, output_size=size))
... input_size = size
... @tf.Module.with_name_scope
... def __call__(self, x):
... for layer in self.layers:
... x = layer(x)
... return x
@tf.Module.with_name_scope
def __call__(self, x):
for layer in self.layers:
x = layer(x)
return x
```
>>> module = MLP(input_size=5, sizes=[5, 5])
>>> module.variables
(<tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
<tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
dtype=float32)>,
<tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
<tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
dtype=float32)>)
"""
# AutoTrackable adds object attributes that users will not expect us to