Обучение с учетом квантования с tf.GradientTape дает ошибку в TensorFlow2.0

Question

Я использую TensorFlow-2.2, tensorflow_model_optimization и Python 3.8. Я пытаюсь квантовать и обучать нейронную сеть LeNet-300-100 Dense, которая содержит разреженность 91,3375%. Это означает, что 91,3375% весов равны нулю. Я следил за учебником Quantization TF и хотел обучить такую ​​разреженную сеть, которая была квантована с помощью tf.GradientTape, а не q_aware_model.fit().

Если вы посмотрите на пример кода, соответствующие фрагменты кода:

quantize_model = tfmot.quantization.keras.quantize_model

# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)


# 'quantize_model' requires recompilation-
q_aware_model.compile(
    optimizer = tf.keras.optimizers.Adam(lr = 0.0012),
    loss=tf.keras.losses.categorical_crossentropy,
    metrics=['accuracy']
)


# Define 'train_one_step()' and 'test_step()' functions here-
@tf.function
def train_one_step(model, mask_model, optimizer, x, y):
    '''
    Function to compute one step of gradient descent optimization
    '''
    with tf.GradientTape() as tape:
        # Make predictions using defined model-
        y_pred = model(x)

        # Compute loss-
        loss = loss_fn(y, y_pred)
        
    # Compute gradients wrt defined loss and weights and biases-
    grads = tape.gradient(loss, model.trainable_variables)
    
    # type(grads)
    # list
    
    # List to hold element-wise multiplication between-
    # computed gradient and masks-
    grad_mask_mul = []
    
    # Perform element-wise multiplication between computed gradients and masks-
    for grad_layer, mask in zip(grads, mask_model.trainable_weights):
        grad_mask_mul.append(tf.math.multiply(grad_layer, mask))
    
    # Apply computed gradients to model's weights and biases-
    optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))

    # Compute accuracy-
    train_loss(loss)
    train_accuracy(y, y_pred)

    return None
    
    
@tf.function
def test_step(model, optimizer, data, labels):
    """
    Function to test model performance
    on testing dataset
    """
    
    predictions = model(data)
    t_loss = loss_fn(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)

    return None



# Train model using 'GradientTape'-
    
# Initialize parameters for Early Stopping manual implementation-
# best_val_loss = 100
# loc_patience = 0
    
for epoch in range(num_epochs):
    
    if loc_patience >= patience:
        print("\n'EarlyStopping' called!\n")
        break
        
    # Reset the metrics at the start of the next epoch
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()
            
    
    for x, y in train_dataset:
        train_one_step(q_aware_model, mask_model, optimizer, x, y)


    for x_t, y_t in test_dataset:
        test_step(q_aware_model, optimizer, x_t, y_t)

    template = 'Epoch {0}, Loss: {1:.4f}, Accuracy: {2:.4f}, Test Loss: {3:.4f}, Test Accuracy: {4:4f}'
    
    '''
    # 'i' is the index for number of pruning rounds-
    history_main[i]['accuracy'][epoch] = train_accuracy.result() * 100
    history_main[i]['loss'][epoch] = train_loss.result()
    history_main[i]['val_loss'][epoch] = test_loss.result()
    history_main[i]['val_accuracy'][epoch] = test_accuracy.result() * 100
    ''' 

    print(template.format(
        epoch + 1, train_loss.result(),
        train_accuracy.result()*100, test_loss.result(),
        test_accuracy.result()*100)
         )
    
    # Count number of non-zero parameters in each layer and in total-
    # print("layer-wise manner model, number of nonzero parameters in each layer are: \n")
    model_sum_params = 0
    
    for layer in winning_ticket_model.trainable_weights:
        # print(tf.math.count_nonzero(layer, axis = None).numpy())
        model_sum_params += tf.math.count_nonzero(layer, axis = None).numpy()
    
    print("Total number of trainable parameters = {0}\n".format(model_sum_params))

    
    # Code for manual Early Stopping:
    if np.abs(test_loss.result() < best_val_loss) >= minimum_delta:
        # update 'best_val_loss' variable to lowest loss encountered so far-
        best_val_loss = test_loss.result()
        
        # reset 'loc_patience' variable-
        loc_patience = 0
        
    else:  # there is no improvement in monitored metric 'val_loss'
        loc_patience += 1  # number of epochs without any improvement

Выдает следующую ошибку:

--------------------------------------------------------------------------- InvalidArgumentError Traceback (последний вызов последний) через 19 20 для x, y в train_dataset: ---> 21 train_one_step(q_aware_model, маска_модель, оптимизатор, x, y)2223

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / def_function.py in call(self, * args, ** kwds) 578 xla_context.Exit () 579 else: -> 580 result = self._call(*args, **kwds) 581 582, если tracing_count == self._get_tracing_count ():

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / def_function.py in _call (self, * args, ** kwds) 642 # Подъем прошел успешно, переменные инициализированы, и мы можем запустить 643 # функция без сохранения состояния. -> 644 return self._stateless_fn (* args, ** kwds) 645 else: 646 canon_args, canon_kwds = \

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / function.py в вызове(self, *args, **kwargs) 2418 с self._lock:
2419 graph_function, args, kwargs =self._maybe_define_function(args, kwargs) -> 2420 return graph_function._filtered_call(args, kwargs) # pylint: disable= protected-access 2421 2422 @property

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / function.py в _filtered_call(self, args, kwargs) 1659 args а такжеkwargs. 1660 "" "-> 1661 return self._call_flat (1662 (t for t in nest.flatten((args, kwargs), expand_composites=True) 1663
if isinstance (t, (ops.Tensor,

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / function.py в _call_flat(self, args, capture_inputs, cancellation_manager)
1743 и executeing_eagerly): 1744 # Лента не просматривается; перейти к запуску функции. -> 1745 return self._build_call_outputs(self._inference_function.call(1746
ctx, args, cancellation_manager=cancellation_manager)) 1747
forward_backward = self._select_forward_and_backward_functions(

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / function.py в вызове (self, ctx, args, cancellation_manager) 591 с _InterpolateFunctionError(self): 592, если cancellation_manager имеет значение None: -> 593 вывода = execute.execute(594 str (self.signature.name),595 num_outputs =self._num_outputs,

~ /.local / lib / python3.8 / site-packages / tensorflow / python / eager / execute.py в quick_execute(op_name, num_outputs, inputs, attrs, ctx, name) 57 try: 58 ctx.ensure_initialized() --> 59 тензоров = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name, 60 input, attrs, num_outputs) 61 кроме core._NotOkStatusException как e:

InvalidArgumentError: var и grad имеют разные формы [10] [100,10] [[узел Adam/Adam/update_4/ResourceApplyAdam (определен на:29) ]] [Op:__inference_train_one_step_20360]

Ошибки могли быть вызваны операцией ввода. Операции с источником ввода, подключенные к узлу Adam/Adam/update_4/ResourceApplyAdam: Mul_4 (определено в:26)
последовательный / Quant_dense_2/BiasAdd/ReadVariableOp/resource (определено в /home/arjun/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper.py:162)

Стек вызовов функций: train_one_step

Есть ли способ объединить квантование модели TF вместе с tf.GradientTape?

Благодарность!