ParamDev's picture
Upload folder using huggingface_hub
a01ef8c verified
#
# Copyright (c) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import shutil
import yaml
import tensorflow as tf
import numpy as np
import time
from tlt.datasets import dataset_factory
from tlt.models import model_factory
from tlt.utils.types import FrameworkType
from PIL import Image
IMAGE_SIZE = 224
def collect_class_labels(dataset_dir):
dataset = dataset_factory.load_dataset(dataset_dir=dataset_dir,
use_case='image_classification',
framework='tensorflow')
return dataset.class_names
def quantize_model(output_dir, saved_model_dir, model):
clean_output_folder(output_dir, 'quantized_models')
quantization_output_dir = os.path.join(output_dir, 'quantized_models',
"vision",
os.path.basename(saved_model_dir))
# Create a tuning workspace directory for INC
root_folder = os.path.dirname(os.path.abspath(__file__))
inc_config_file = os.path.join(root_folder, "config.yaml")
# inc_config_file = 'vision/config.yaml'
model.quantize(quantization_output_dir, inc_config_file)
def clean_output_folder(output_dir, model_name):
folder_path = os.path.join(output_dir, model_name)
if os.path.exists(folder_path):
shutil.rmtree(os.path.join(output_dir, model_name))
def train_vision_wl(dataset_dir, output_dir, model="resnet_v1_50",
batch_size=32,
epochs=5, save_model=True, quantization=False, bf16=True):
# Clean the output folder first
clean_output_folder(output_dir, model)
dict_metrics = {}
# Loading the model
tstart = time.time()
model = model_factory.get_model(model_name=model,
framework=FrameworkType.TENSORFLOW)
tend = time.time()
print("\nModel Loading time (s): ", tend - tstart)
# Load the dataset from the custom dataset path
# Data loading and preprocessing #
dataset = dataset_factory.load_dataset(dataset_dir=dataset_dir,
use_case='image_classification',
framework='tensorflow',
shuffle_files=True)
print("Class names:", str(dataset.class_names))
dataset.preprocess(model.image_size, batch_size=batch_size,
add_aug=['hvflip', 'rotate'])
dataset.shuffle_split(train_pct=.80, val_pct=.20)
# Finetuning #
tstart = time.time()
history = model.train(dataset, output_dir=output_dir, epochs=epochs,
seed=10,
enable_auto_mixed_precision=bf16,
extra_layers=[1024, 512])
tend = time.time()
print("\nTotal Vision Finetuning time (s): ", tend - tstart)
dict_metrics['e2e_training_time'] = tend - tstart
metrics = model.evaluate(dataset)
for metric_name, metric_value in zip(model._model.metrics_names, metrics):
print("{}: {}".format(metric_name, metric_value))
dict_metrics[metric_name] = metric_value
print('dict_metrics:', dict_metrics)
print('Finished Fine-tuning the vision model...')
if save_model:
saved_model_dir = model.export(output_dir)
if quantization:
print('Quantizing the model')
quantize_model(output_dir, saved_model_dir, model)
print("Done finetuning the vision model ............")
return (model, history, dict_metrics)
def infer_vision_wl(model, image_location):
image_shape = (model.image_size, model.image_size)
image = Image.open(image_location).resize(image_shape)
# Get the image as a np array and call predict while adding a batch
# dimension (with np.newaxis)
image = np.array(image)/255.0
result = model.predict(image[np.newaxis, ...], 'probabilities')[0]
return result
def infer_int8_vision_wl(model, image_location):
image_shape = (IMAGE_SIZE, IMAGE_SIZE)
image = Image.open(image_location).resize(image_shape)
# Get the image as a np array and call predict while
# adding a batch dimension (with np.newaxis)
image = np.array(image)/255.0
image = image[np.newaxis, ...].astype('float32')
infer = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
# result = model.predict(image[np.newaxis, ...])
# result=model.predict(image[np.newaxis, ...], 'probabilities')[0]
output_name = list(infer.structured_outputs.keys())
result = infer(tf.constant(image))[output_name[0]][0]
return result
def preprocess_dataset(dataset_dir, image_size, batch_size):
"""
Load and preprocess dataset
"""
dataset = dataset_factory.load_dataset(dataset_dir=dataset_dir,
use_case='image_classification',
framework='tensorflow',
shuffle_files=False)
dataset.preprocess(image_size, batch_size)
class_dict = reverse_map(dataset.class_names)
return dataset, class_dict
def reverse_map(class_names):
class_dict = {}
i = 0
for c in class_names:
class_dict[i] = c
i = i + 1
return class_dict
def load_model(model_name, saved_model_dir):
vision_model = model_factory.load_model(model_name, saved_model_dir,
"tensorflow",
"image_classification")
return vision_model
def run_inference_per_patient(model, patient_dict,class_names):
results = {}
class_dict = reverse_map(class_names)
for key, value in patient_dict.items():
print(key, '->', value)
results[key] = {}
for image in value:
pred_prob = infer_vision_wl(model,image).numpy().tolist()
infer_result_patient = [
{
"label": image.split('/')[-2],
"pred": class_dict[np.argmax(pred_prob).tolist()],
"pred_prob": pred_prob
}
]
results[key][image.split('/')[-1]] = infer_result_patient
print(results)
return results
def run_inference(test_data_dir, saved_model_dir, class_labels,
model_name="resnet_v1_50", vision_int8_inference=False,
report="output.yaml"):
# Load the vision model
tstart = time.time()
vision_model_dir = saved_model_dir
test_dir = test_data_dir
labels = class_labels
predictions_report_save_file = report
predictions_report = {}
predictions_report["metric"] = {}
predictions_report["results"] = {}
# Load model
vision_model = model_factory.load_model(model_name, vision_model_dir,
"tensorflow",
"image_classification")
if vision_int8_inference:
vision_int8_model = tf.saved_model.load(vision_model_dir)
tend = time.time()
print("\n Vision Model Loading time: ", tend - tstart)
# Load dataset for metric evaluation
dataset, class_dict = preprocess_dataset(test_data_dir,
vision_model.image_size, 32)
metrics = vision_model.evaluate(dataset)
for metric_name, metric_value in zip(vision_model._model.metrics_names,
metrics):
print("{}: {}".format(metric_name, metric_value))
predictions_report["metric"][metric_name] = metric_value
tstart = time.time()
for label in os.listdir(test_dir):
print("Infering data in folder: ", label)
fns = os.listdir(os.path.join(test_dir, label))
for fn in fns:
patient_id = fn
fn = os.path.join(os.path.join(test_dir, label, fn))
# ------------------------
# call inference on vision WL
# ------------------------
if vision_int8_inference:
result_vision = infer_int8_vision_wl(vision_int8_model, fn)
else:
result_vision = infer_vision_wl(vision_model, fn)
pred_prob = result_vision.numpy().tolist()
infer_result_patient = [
{
"label": label,
"pred": class_dict[np.argmax(pred_prob).tolist()],
"pred_prob": pred_prob
}
]
predictions_report["label"] = labels
predictions_report["label_id"] = list(class_dict.keys())
predictions_report["results"][patient_id] = infer_result_patient
with open(predictions_report_save_file, 'w') as file:
_ = yaml.dump(predictions_report, file, )
print("Vision inference time: ", time.time() - tstart)