Может быть, вы могли бы создать подкласс плотного слоя? Что-то вроде
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
(Опять не слишком тщательно протестировано и определенно нуждается в некоторой полировке, но я думаю,идея должна работать)
Edit 3: поиск в Google дал мне то, что я изначально не смог найти: https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/sparsity/keras migth предоставляет инструменты для более простого построения обрезанной модели.
Edit4 (дальнейшее объяснение роли _kernel2 и _bias2):
Для простоты я сделаю объяснение без предвзятости, но mutatis mutandis все работает так же с уклоном. Предположим, что входные данные для вашего плотного слоя n-мерные, а ваши выходные данные m-мерные, тогда плотный слой умножает входные данные на матрицу m-на-n, которую мы кратко будем называть K (это ядро),
Обычно мы хотим узнать правильные записи K с помощью некоторого градиентного метода оптимизации, но в вашем случае вы хотите сохранить некоторые записи фиксированными. Вот почему в этом пользовательском плотном слое мы разделяем K следующим образом:
K = T * K1 + (1 - T) * K2,
, где
- T - это матрица m-by-n, состоящая из 0 и 1,
- , звездочка обозначает поэлементное умножение
- 1 - матрица m-by-n с 1 для каждой записи
- K1 - это матрица m-by-n, которую можно выучить
- K2 - это матрица m-by-n, которая фиксируется (постоянная) во время тренировки.
Еслимы смотрим на записи K, тогда K [i, j] = T [i, j] * K1 [i, j] + (1-T [i, j]) * K2 [i, j] = K1 [i, j] если T [i, j] == 1, то еще K2 [i, j]. Поскольку в последнем случае значение K1 [i, j] не влияет на результат умножения на K, его градиент равен 0 и не должен изменяться (и даже если он действительно изменяется из-за числовых ошибок, это не должновлияют на значение K [i, j]).
По существу, записи K [i, j] для K, для которых T [i, j] == 0, являются фиксированными (со значением, сохраненным в K2), и те, для которых T [i, j]== 1 можно обучить.