如何在TensorFlow中锁定Tensor的特定值?

问题描述 投票:2回答:1

我正在尝试将彩票假设应用于TensorFlow 2.0(使用Keras接口)编写的简单神经网络,如下所示:

net = models.Sequential()
net.add(layers.Dense(256, activation="softsign", name="Dense0", bias_initializer="ones"))
net.add(layers.Dense(128, activation="softsign", name="Dense1", bias_initializer="ones"))
net.add(layers.Dense(64, activation="softsign", name="Dense2", bias_initializer="ones"))
net.add(layers.Dense(32, activation="softsign", name="Dense3", bias_initializer="ones"))
net.add(layers.Dense(1, activation="tanh", name="Output", bias_initializer="ones"))

然后我使用Adam优化器和二进制交叉熵损失训练我的网络:

net.compile(optimizer=optimizers.Adam(learning_rate=0.001),
            loss=losses.BinaryCrossentropy(), metrics=["accuracy"])
net.fit(x_train, y_train, epochs=10, batch_size=32, validation_data=(x_test, y_test))

经过培训,我想在网络中锁定特定权重。问题是,我只能使用tensorflow.Variable(..., trainable=False)将张量锁定为不可训练(据我所知,fas),但是这样做是将图形的整个节点设置为不可训练,并且我只想特定边缘。我可以使用以下代码遍历网络的所有Tensor实例:

for i in range(len(net.layers)):
    for j in range(net.layers[i].variables[0].shape[0]):
        for k in range(net.layers[i].variables[0][j].shape[0]):
            ...

但是我不知道接下来要做什么。有人知道我可以这样做的简单方法吗?

python tensorflow keras tensorflow2.0
1个回答
0
投票

也许您可以继承Dense层的子类?有点像

class PrunableDense(keras.layers.Dense):


    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.deleted_channels = None
        self.deleted_bias = None
        self._kernel=None
        self._bias=None


    def build(self, input_shape):
        last_dim = input_shape[-1]
        self._kernel = self.add_weight(
            'kernel',
            shape=[last_dim, self.units],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
            dtype=self.dtype,
            trainable=True)
        self.deleted_channels = tf.ones([last_dim, self.units]) # we'll use this to prune the network
        if self.use_bias:
            self._bias = self.add_weight(
                'bias',
                shape=[self.units,],
                initializer=self.bias_initializer,
                regularizer=self.bias_regularizer,
                constraint=self.bias_constraint,
                dtype=self.dtype,
                trainable=True)
            self.deleted_bias = tf.ones([self.units,])

    @property
    def kernel(self):
        """gets called whenever self.kernel is used"""
        # only the weights that haven't been deleted should be non-zero
        # deleted weights are 0.'s in self.deleted_channels
        return self.deleted_channels * self._kernel  

    @property
    def bias(self):
        #similar to kernel
        if not self.use_bias:
            return None
        else:
            return self.deleted_bias * self._bias

    def prune_kernel(self, to_be_deleted):
        """
        Delete some channels
        to_be_deleted should be a tensor or numpy array of shape kernel.shape
        containing 1's at the locations where weights should be kept, and 0's 
        at the locations where weights should be deleted.
        """
        self.deleted_channels *= to_be_deleted

    def prune_bias(self, to_be_deleted):
        assert(self.use_bias)
        self.deleted_bias *= to_be_deleted

    def prune_kernel_below_threshold(self, threshold=0.01):
        to_be_deleted = tf.cast(tf.greater(self.kernel, threshold), tf.float32)
        self.deleted_channels *= to_be_deleted

    def prune_bias_below_threshold(self, threshold=0.01):
        assert(self.use_bias)
        to_be_deleted = tf.cast(tf.greater(self.bias, threshold), tf.float32)
        self.deleted_bias *= to_be_deleted

我还没有对它进行过全面的测试,它肯定需要改进,但是我认为这个想法应该可行。

编辑:我在上面写的是假设您想像彩票假设中那样修剪网络,但是如果您只想冻结部分权重,则可以执行类似的操作,但是要添加带有非零条目的Frozen_kernel属性仅在self.deleted_channels为0的地方,并将其添加到可训练的内核中。

编辑2:在先前的编辑中,我的意思如下:

class FreezableDense(keras.layers.Dense):


    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.trainable_channels = None
        self.trainable_bias = None
        self._kernel1 = None
        self._bias1 = None
        self._kernel2 = None
        self._bias2 = None


    def build(self, input_shape):
        last_dim = input_shape[-1]
        self._kernel1 = self.add_weight(
            'kernel1',
            shape=[last_dim, self.units],
            initializer=self.kernel_initializer,
            regularizer=self.kernel_regularizer,
            constraint=self.kernel_constraint,
            dtype=self.dtype,
            trainable=True)
        self._kernel2 = tf.zeros([last_dim, self.units])
        self.trainable_channels = tf.ones([last_dim, self.units]) # we'll use this to freeze parts of the network
        if self.use_bias:
            self._bias1 = self.add_weight(
                'bias',
                shape=[self.units,],
                initializer=self.bias_initializer,
                regularizer=self.bias_regularizer,
                constraint=self.bias_constraint,
                dtype=self.dtype,
                trainable=True)
            self._bias2 = tf.zeros([self.units,])
            self.trainable_bias = tf.ones([self.units,])

    @property
    def kernel(self):
        """gets called whenever self.kernel is used"""
        # frozen 
        return self.trainable_channels * self._kernel1 + (1-self.trainable_channels * self._kernel2) 

    @property
    def bias(self):
        #similar to kernel
        if not self.use_bias:
            return None
        else:
            return self.trainable_bias * self._bias1 + (1-self.trainable_bias)*self._bias2

    def freeze_kernel(self, to_be_frozen):
        """
        freeze some channels
        to_be_frozen should be a tensor or numpy array of shape kernel.shape
        containing 1's at the locations where weights should be kept trainable, and 0's 
        at the locations where weights should be frozen.
        """
        # we want to do two things: update the weights in self._kernel2 
        # and update self.trainable_channels
        # first we update self._kernel2 with all newly frozen weights
        newly_frozen = 1 - tf.maximum((1 - to_be_frozen) - (1 - self.trainable_channels), 0)
        # the above should have 0 only where to_be_frozen is 0 and self.trainable_channels is 1
        # if I'm not mistaken that is
        newly_frozen_weights = (1-newly_frozen)*self._kernel1
        self._kernel2 += newly_frozen_weights

        # now we update self.trainable_channels:
        self.trainable_channels *= to_be_frozen

    def prune_bias(self, to_be_deleted):
        assert(self.use_bias)
        newly_frozen = 1 - tf.maximum((1 - to_be_frozen) - (1 - self.trainable_bias), 0)
        newly_frozen_bias = (1-newly_frozen)*self._bias1
        self._bias2 += newly_frozen_bias
        self.trainable_bias *= to_be_frozen

(再次没有经过详尽的测试,绝对需要完善,但我认为这个想法应该可行)

编辑3:使用Google进行更多的搜索时发现了我最初找不到的东西:https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/sparsity/keras migth提供了更轻松地构建修剪模型的工具。

© www.soinside.com 2019 - 2024. All rights reserved.