Efficiently add noise to all trainable weights in a model

user_11714396 0
Viewed 737

How can I efficiently access all trainable variables of a Keras model within a tf.fuction to add custom noise to all variables?

Let's assume this simple model:

my_model = Sequential()
my_model.add(Dense(300, input_dim=40, activation='relu'))
my_model.add(Dense(200, activation='relu'))
my_model.add(Dense(8, activation='sigmoid'))

Running eagerly, I can do this by:

@tf.function
def weight_perturbation(model, generator):
    n_layers = len(model.layers)

    # iterate over all layers
    for i in tf.range(n_layers):
        trainable_weights = model.layers[i].trainable_variables

        # iterate over all weight vectors in a layer 
        for j in tf.range(len(trainable_weights)):
            trainable_weights[j].assign_add(generator.normal(tf.shape(trainable_weights[j]), dtype=tf.float64))

However, when not running eagerly, I get the following error:

     trainable_weights = model.layers[i].trainable_variables

TypeError: list indices must be integers or slices, not Tensor

How can I circumvent this problem? I had a look at tf.gather but this does not seem to work with lists.

The best would be, if there is a way to get all trainable weights of a model as a flattened tensor without looping. Unfortunately, I haven't found anything like this.

tensorflow2.0 keras tensorflow python
1 Answers

It works for me if I iterate through the layers and weights rather than iterating through tf.range:

def weight_perturbation(model):
    for layer in model.layers:
        trainable_weights = layer.trainable_variables

        for weight in trainable_weights :
            random_weights = tf.random.uniform(tf.shape(weight), 
                                               1e-4, 1e-5, dtype=tf.float32)
            weight.assign_add(random_weights)

I changed your data type to tf.float32, as it should be in most cases. I added the weight manipulation here:

import tensorflow as tf
from tensorflow import keras as K
from tensorflow.keras.layers import Conv2D, Flatten, Dense, MaxPooling2D, Dropout
from tensorflow import nn as nn
from functools import partial

(xtrain, ytrain), (xtest, ytest) = tf.keras.datasets.mnist.load_data()

train = tf.data.Dataset.from_tensor_slices((xtrain, ytrain))
test = tf.data.Dataset.from_tensor_slices((xtest, ytest))

def prepare(inputs, outputs):
    inputs = tf.cast(inputs, tf.float32)
    inputs = tf.divide(x=inputs, y=255)
    inputs = tf.expand_dims(inputs, -1)
    targets = tf.one_hot(outputs, depth=10)
    return inputs, targets

train = train.map(prepare).batch(64)
test = test.map(prepare).batch(64)

class MyCNN(K.Model):
    def __init__(self):
        super(MyCNN, self).__init__()
        Conv = partial(Conv2D, kernel_size=(3, 3), activation=nn.relu)
        MaxPool = partial(MaxPooling2D, pool_size=(2, 2))

        self.conv1 = Conv(filters=8)
        self.maxp1 = MaxPool()
        self.conv2 = Conv(filters=8)
        self.maxp2 = MaxPool()
        self.flatt = Flatten()
        self.dens1 = Dense(8, activation=nn.relu)
        self.drop1 = Dropout(.5)
        self.dens2 = Dense(10, activation=nn.softmax)

    def call(self, x, training=None, **kwargs):
        x = self.conv1(x)
        x = self.maxp1(x)
        x = self.conv2(x)
        x = self.maxp2(x)
        x = self.flatt(x)
        x = self.dens1(x)
        x = self.drop1(x)
        x = self.dens2(x)
        return x

model = MyCNN()

loss_object = tf.losses.CategoricalCrossentropy(from_logits=False)

def compute_loss(model, x, y, training):
  out = model(inputs=x, training=training)
  loss = loss_object(y_true=y, y_pred=out)
  return loss, out

def get_grad(model, x, y):
    with tf.GradientTape() as tape:
        loss, out = compute_loss(model, x, y, training=False)
    return loss, tape.gradient(loss, model.trainable_variables), out

def weight_perturbation(model):
    for layer in model.layers:
        trainable_weights = layer.trainable_variables

        for weight in trainable_weights :
            random_weights = tf.random.uniform(tf.shape(weight),
                                               1e-4, 1e-5, dtype=tf.float32)
            weight.assign_add(random_weights)

optimizer = tf.optimizers.Adam()

verbose = "Epoch {:2d} Loss: {:.3f} TLoss: {:.3f} Acc: {:.3%} TAcc: {:.3%}"

for epoch in range(1, 10 + 1):
    train_loss = tf.metrics.Mean()
    train_acc = tf.metrics.CategoricalAccuracy()
    test_loss = tf.metrics.Mean()
    test_acc = tf.metrics.CategoricalAccuracy()

    weight_perturbation(model)

    for x, y in train:
        loss_value, grads, out = get_grad(model, x, y)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        train_loss.update_state(loss_value)
        train_acc.update_state(y, out)

    for x, y in test:
        loss_value, _, out = get_grad(model, x, y)
        test_loss.update_state(loss_value)
        test_acc.update_state(y, out)

    print(verbose.format(epoch,
                         train_loss.result(),
                         test_loss.result(),
                         train_acc.result(),
                         test_acc.result()))
user_109083750
Related