Source code for akida_models.extract

#!/usr/bin/env python
# ******************************************************************************
# Copyright 2023 Brainchip Holdings Ltd.
#
# 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.
# ******************************************************************************
"""Sample extraction from datasets"""

import numpy as np
import tensorflow as tf
from typing import Iterable
import itertools


[docs]def extract_samples(out_file, dataset, nb_samples=1024): """Extracts samples from dataset and save them to a npz file. Args: out_file (str): name of output file dataset (numpy.ndarray or tf.data.Dataset): dataset for extract samples nb_samples (int, optional): number of samples. Defaults to 1024. """ # The expected number of samples if isinstance(dataset, np.ndarray): if len(dataset) < nb_samples: raise ValueError("Not enough samples in the dataset.") samples_x = dataset[0:nb_samples].astype('float32') elif isinstance(dataset, tf.data.Dataset): samples_x, _ = next(dataset.as_numpy_iterator()) elif isinstance(dataset, Iterable): dataset = itertools.islice(dataset, nb_samples) # drop labels xs, _ = itertools.tee(dataset) xs = (x[0] for x in xs) samples_x = next(xs) else: raise ValueError("Dataset format not supported.") samples = {"data": samples_x} np.savez(out_file, **samples) print(f"Samples saved as {out_file}")