Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

`tf.split` or `tf.transpose` cause errors for quantize-aware training with `quantize_apply`

Abierto
  • ¿Algo no está bien?
#1,062 4 comentarios 0 reacciones 1 asignado Ver en GitHub

@cdh4696 ya está trabajando en esto.

Desde el 2/5/2023.

Evaluación

Este issue todavía no se ha evaluado.

Descripción

bug

Describe the bug

We are trying to implement some network like ShuffleNetV2 but encounter some error when quantize_apply the model.

image

I believe ShuffleNet or related ideas are popular in edge devices, please kindly help us to resolve this proble.

Any advice is welcome.

System information

TensorFlow version (installed from source or binary): 2.7.0

TensorFlow Model Optimization version (installed from source or binary): 0.7.0

Python version: 3.8.13

Describe the expected behavior

Just add quantization-aware operator in to the model.

Describe the current behavior

When running the provided code, either the tf.transpose or tf.split will cause error to Tensorflow Model Optimization.

The error message due to tf.split before convolution layers:

ValueError: Exception encountered when calling layer "bn3" (type BatchNormalization).

Shape must be rank 4 but is rank 5 for '{{node bn3/FusedBatchNormV3}} = FusedBatchNormV3[T=DT_FLOAT, U=DT_FLOAT, data_format="NHWC", epsilon=0.001, exponential_avg_factor=1, is_training=false](Placeholder, bn3/ReadVariableOp, bn3/ReadVariableOp_1, bn3/FusedBatchNormV3/ReadVariableOp, bn3/FusedBatchNormV3/ReadVariableOp_1)' with input shapes: [1,?,128,128,32], [32], [32], [32], [32].

The error message due to tf.transpose:

ValueError: Exception encountered when calling layer "tf.compat.v1.transpose" (type TFOpLambda).

Dimension must be 6 but is 5 for '{{node tf.compat.v1.transpose/transpose}} = Transpose[T=DT_FLOAT, Tperm=DT_INT32](tf.compat.v1.transpose/transpose/a, tf.compat.v1.transpose/transpose/perm)' with input shapes: [1,?,128,128,2,32], [5].

Code to reproduce the issue

Just run the following code you will get the error message due to tf.split.

from __future__ import annotations

from typing import Callable, Optional

import tensorflow as tf
import tensorflow_model_optimization as tfmot
from tensorflow.keras import layers


SKIP_LAYER = [
    "resize",
    "Resize",
    "reshape",
    "Reshape",
    "concat",
    "Concat" "ExpandDims",
    "Repeats",
    "Shape",
    "strided_slice",
    "Tile",
]


def quantize_model(
    model: tf.keras.Model,
    annotate: Optional[Callable] = None,
    quantize_scope: Optional[dict[str, tf.keras.layers.Layer]] = None,
) -> tf.keras.Model:
    quantize_scope = {} if quantize_scope is None else quantize_scope

    def annotate(layer):
        if any([name in layer.name for name in SKIP_LAYER]):
            return layer
        else:
            return tfmot.quantization.keras.quantize_annotate_layer(layer)

    anno_model = tf.keras.models.clone_model(model, clone_function=annotate)
    with tfmot.quantization.keras.quantize_scope(quantize_scope):
        model = tfmot.quantization.keras.quantize_apply(anno_model)

    return model


def channel_shuffle(tensor: tf.Tensor, groups: int = 2) -> tf.Tensor:
    """Channel shuffle operation."""
    _, height, width, num_channels = tensor.shape.as_list()
    assert num_channels % groups == 0

    tensor = tf.reshape(tensor, [-1, height, width, groups, num_channels // groups])
    tensor = tf.transpose(tensor, [0, 1, 2, 4, 3])
    tensor = tf.identity(tensor, name="channel_shuffle")

    tensor = tf.reshape(tensor, [-1, height, width, num_channels])
    return tensor


def simple_nn(img_input: tf.Tensor) -> tf.Tensor:
    latent = layers.Conv2D(32, 1, padding="same", use_bias=False, name="conv1")(img_input)
    latent = layers.BatchNormalization(name="bn1")(latent)
    latent = layers.ReLU(name="relu1")(latent)

    latent = layers.DepthwiseConv2D(3, 1, padding="same", name="conv2")(img_input)
    latent = layers.BatchNormalization(name="bn2")(latent)

    latent = layers.Conv2D(32, 1, padding="same", use_bias=False, name="conv3")(img_input)
    latent = layers.BatchNormalization(name="bn3")(latent)
    latent = layers.ReLU(name="relu3")(latent)

    return latent


def split_like_nn(img_input: tf.Tensor) -> tf.Tensor:
    latent = layers.Conv2D(64, 1, padding="same", use_bias=False, name="conv0")(img_input)
    latent = layers.BatchNormalization(name="bn0")(latent)
    latent = layers.ReLU(name="relu0")(latent)

    latent_0, latent_1 = tf.split(latent, 2, axis=-1)
    latent_0 = simple_nn(latent_0)
    latent = tf.concat([latent_0, latent_1], axis=-1)

    latent = channel_shuffle(latent)

    return latent


if __name__ == "__main__":
    img_input = tf.keras.Input((128, 128, 1), dtype=tf.float32, name="img")

    outputs = split_like_nn(img_input)

    model = tf.keras.Model(inputs=img_input, outputs=outputs, name="PoseNetV2")
    model.summary()

    model_qat = quantize_model(model)
    model_qat.summary()

You can just comment the following three lines of code will get the error message from tf.transpose.

 latent_0, latent_1 = tf.split(latent, 2, axis=-1)
 latent_0 = simple_nn(latent_0)
 latent = tf.concat([latent_0, latent_1], axis=-1)
Lenguaje dominante
Python
Estrellas
1.6k
Forks
349
Sin PR fusionados en 30 d

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de tensorflow/model-optimization

Todos los issues de tensorflow/model-optimization

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.