这是另一种基于二进制掩码的解决方案。
"""Solution based on binary mask.
- We just add this mask to inputs, instead of multiplying."""
class AddToEven(tf.keras.Model):
def __init__(self):
super(AddToEven, self).__init__()
def build(self, inputshape):
self.built = True # Actually nothing to build with, becuase we don't have any variables or weights here.
@tf.function
def call(self, inputs):
w = inputs.get_shape()[-1]
# 1-d mask generation for w-axis (activate even indices only)
m_w = tf.range(w) # [0, 1, 2,... w-1]
m_w = ((m_w%2)==0) # [True, False, True ,...] with dtype=tf.bool
# Apply 1-d mask to 2-d input
m_w = tf.expand_dims(m_w, axis=0) # just extend dimension as to be (1, W)
m_w = tf.cast(m_w, dtype=inputs.dtype) # in advance, we need to convert dtype
# Here, we just add this (1, W) mask to (H,W) input magically.
outputs = inputs + m_w # This add operation is allowed in both TF and numpy!
return tf.reshape(outputs, inputs.get_shape())
在这里进行健全性检查。
# sanity-check as model
model = AddToEven()
model.build(tf.TensorShape([None, None]))
z = model(tf.zeros([2,4]))
print(z)
结果(使用 TF 2.1)是这样的。
tf.Tensor(
[[1. 0. 1. 0.]
[1. 0. 1. 0.]], shape=(2, 4), dtype=float32)
-------- 以下是之前的回答 --------
您需要在 build() 方法中创建 tf.Variable 。
它还允许通过 shape=(None,) 动态调整大小。
在下面的代码中,我将输入形状指定为(无,无)。
class AddToEven(tf.keras.Model):
def __init__(self):
super(AddToEven, self).__init__()
def build(self, inputshape):
self.v = tf.Variable(initial_value=tf.zeros((0,0)), shape=(None, None), trainable=False, dtype=tf.float32)
@tf.function
def call(self, inputs):
self.v.assign(inputs)
self.v[:, ::2].assign(self.v[:, ::2] + 1)
return self.v.value()
我用 TF 2.1.0 和 TF1.15 测试了这段代码
# test
add_to_even = AddToEven()
z = add_to_even(tf.zeros((2,4)))
print(z)
Result:
tf.Tensor(
[[1. 0. 1. 0.]
[1. 0. 1. 0.]], shape=(2, 4), dtype=float32)
附:还有其他一些方法,例如使用 tf.numpy_function(),或生成掩码函数。