TensorFlow2 推理

参考：migrating_checkpoints

下面以模型 resnet_v2_50 为例展示。

需要克隆项目 models，然后执行如下操作。

import tensorflow as tf
try:
    tf1 = tf.compat.v1
except (ImportError, AttributeError):
    tf1 = tf
tf.get_logger().setLevel('ERROR')

切换到 models/research/slim 目录下：

%cd /media/pc/data/lxw/ai/tasks/models/research/slim

%%bash 
pip install --upgrade tf_slim tf-keras -i https://pypi.tuna.tsinghua.edu.cn/simple

将 TF1 升级为 TF2：

from nets import resnet_v2
import tf_slim as slim


class ResnetV2_50(tf.Module):
    @tf.function(input_signature=[tf.TensorSpec([None, 299, 299, 3], tf.float32, name="data")])
    @tf1.keras.utils.track_tf1_style_variables
    def __call__(self, data):
        with slim.arg_scope(resnet_v2.resnet_arg_scope()):
            logits, end_points = resnet_v2.resnet_v2_50(
                data, 
                num_classes=1001,
                global_pool=True,
                is_training=False,
                scope="resnet_v2_50"
            )
        del end_points
        return tf.nn.softmax(logits) 

预处理：

from PIL import Image
import numpy as np
from nets import resnet_v2
from tvm_book.data.classification import ImageFolderDataset
import tf_slim as slim
import tensorflow as tf


@tf.function
def preprocessing(
    image,
    use_grayscale=False,
    central_fraction=0.875,
    central_crop=True,
    height=299,
    width=299,
    mean: tuple[float, ...] = (0.485, 0.456, 0.406),
    std: tuple[float, ...] = (1, 1, 1)
):
    # image = tf.constant(image)
    if image.dtype != tf.float32:
        image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    if use_grayscale:
        image = tf.image.rgb_to_grayscale(image)
    if central_crop and central_fraction:
        image = tf.image.central_crop(image, central_fraction=central_fraction)
    if height and width:
        image = tf.expand_dims(image, 0)
        image = tf.image.resize(image, [height, width],
                                method='bilinear',
                                preserve_aspect_ratio=False,
                                antialias=False)
        image = tf.squeeze(image, [0])
    image = tf.subtract(image, mean)
    image = tf.divide(image, std)
    return image


# 预处理
root = "/media/pc/data/lxw/home/data/datasets/ILSVRC/val"
valset = ImageFolderDataset(root)
image, label_id = valset[1001]
model_dir = 'temp/resnet_v2_50'
# remove_dir(model_dir)
processed_image = preprocessing(
    image,
    use_grayscale=False,
    central_fraction=0.875,
    central_crop=True,
    height=299,
    width=299,
    mean=(0.485, 0.456, 0.406),
    std=(1, 1, 1)
)
np_processed_images = np.expand_dims(processed_image.numpy(), axis=0)

前向推理：

model = ResnetV2_50()
model(tf.ones(shape=(1, 299, 299, 3), dtype=tf.float32))
ckpt = tf.train.Checkpoint(model=model)
ckpt.restore(".temp/checkpoints/resnet_v2_50.ckpt") # 更新模型参数
outputs = model(np_processed_images)
outputs = outputs.numpy()

打印标签信息：

from tvm_book.data.imagenet.classification import ImageNet1kAttr

imagenet1k_attr = ImageNet1kAttr()
sorted_inds = outputs[0].argsort()[::-1]
topk = 5
print(f"真实标签：{imagenet1k_attr.classes_long[label_id]}")
for sorted_ind in sorted_inds[:topk]:
    label = imagenet1k_attr.classes_long[sorted_ind-1]
    print(f"{sorted_ind-1}: {label.ljust(38)}\t{outputs[0, sorted_ind]}")

将其模型和参数保存下来：

module_with_signature_path = 'temp/module_with_signature'
call = model.__call__.get_concrete_function(tf.TensorSpec([1, 299, 299, 3], tf.float32, name="data"))
tf.saved_model.save(model, module_with_signature_path, signatures=call)
imported_with_signatures = tf.saved_model.load(module_with_signature_path)
infer = imported_with_signatures.signatures['serving_default']
labeling = infer(tf.constant(np_processed_images))['output_0']
# gdef = model.__call__.get_concrete_function().graph.as_graph_def(add_shapes=True)
# gdef_ops = list(set([n.op for n in gdef.node]))
# gdef = infer.graph.as_graph_def(add_shapes=True)