具有不同数据类型的神经网络层

Question

I would like to make a neural network whose layers consists of different data types. for example, one layer could be of INT16 data type and one layer could be of Float16 or INT32.

I am trying to make it using the dtype keyword. I am able to achieve this for different float data types. ie float16 and float32 using the following code snippet. but when i tried to do the same for INT32 and Float, it gave me an error. I am not able to figure it out.

TypeError: Unable to build Dense layer with non-floating point dtype <dtype: 'int16'>

Could anyone help me out to make such a network with different datatypes

# Code Snippet 
from tensorflow.keras.models import Model
i = Input(shape= x_train[0].shape, dtype=tf.float16)
x = Flatten(dtype=tf.float16)(i)
x = Dense(64, activation='relu',dtype=tf.float32)(x)
x = Dense(64, activation='relu')(x)
x = Dense(64, activation='relu', dtype=tf.float16)(x)
x = Dense(64, activation='relu', dtype=tf.float16)(x)
x = Dense(10, activation='softmax')(x)

model = Model(i, x)

# data type of model's layer can be observed as 

input_14  :  float16
flatten_13  :  float16
dense_57  :  float32
dense_58  :  float32
dense_59  :  float16
dense_60  :  float16
dense_61  :  float32

The code which gave error

# Step 4- Now let's make the ANN!
from tensorflow.keras.models import Model
i = Input(shape= x_train[0].shape, dtype=tf.int16)
x = Flatten(dtype=tf.int16)(i)
x = Dense(64, activation='relu',dtype=tf.int16)(x)        # error occurred at this line
x = Dense(64, activation='relu')(x)
x = Dense(64, activation='relu', dtype=tf.float16)(x)
x = Dense(64, activation='relu', dtype=tf.float16)(x)
x = Dense(10, activation='softmax')(x)

Answer 1

This is not something you can do just by casting to tf.int32 . What you're talking about is called model quantization and is an extremely complicated process. Tensorflow, via Tensorflow Lite, offers many resources for quantization aware training. I suggest you start there.

Quantization aware training emulates inference-time quantization, creating a model that downstream tools will use to produce actually quantized models. The quantized models use lower-precision (eg 8-bit instead of 32-bit float), leading to benefits during deployment.

With this example you can train a MNIST model using integers:

import tempfile
import os
import tensorflow as tf
from tensorflow import keras

mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

train_images = train_images / 255.0
test_images = test_images / 255.0

model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(10)
])

# Train the digit classification model
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(
  train_images,
  train_labels,
  epochs=1,
  validation_split=0.1,
)

import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model

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

# `quantize_model` requires a recompile.
q_aware_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

train_images_subset = train_images[0:1000] # out of 60000
train_labels_subset = train_labels[0:1000]

q_aware_model.fit(train_images_subset, train_labels_subset,
                  batch_size=500, epochs=1, validation_split=0.1)

_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

_, q_aware_model_accuracy = q_aware_model.evaluate(
   test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Quant test accuracy:', q_aware_model_accuracy)