diff options
| author | Nathan Bailey <nathan.bailey@arm.com> | 2024-03-20 08:13:39 +0000 |
|---|---|---|
| committer | Nathan Bailey <nathan.bailey@arm.com> | 2024-03-28 07:17:32 +0000 |
| commit | f3f3ab451968350b8f6df2de7c60b2c2b9320b59 (patch) | |
| tree | 05d56c8e41de9b32f8054019a21b78628151310d /tests/test_nn_tensorflow_optimizations_clustering.py | |
| parent | 5f063ae1cfbfa2568d2858af0a0ccaf192bb1e8d (diff) | |
| download | mlia-f3f3ab451968350b8f6df2de7c60b2c2b9320b59.tar.gz | |
feat: Update Vela version
Updates Vela Version to 3.11.0 and TensorFlow version to 2.15.1
Required keras import to change:
from keras.api._v2 import keras needed instead of calling tf.keras
Subsequently tf.keras.X needed to change to keras.X
Resolves: MLIA-1107
Signed-off-by: Nathan Bailey <nathan.bailey@arm.com>
Change-Id: I53bcaa9cdad58b0e6c311c8c6490393d33cb18bc
Diffstat (limited to 'tests/test_nn_tensorflow_optimizations_clustering.py')
| -rw-r--r-- | tests/test_nn_tensorflow_optimizations_clustering.py | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/tests/test_nn_tensorflow_optimizations_clustering.py b/tests/test_nn_tensorflow_optimizations_clustering.py index 72ade58..11036ad 100644 --- a/tests/test_nn_tensorflow_optimizations_clustering.py +++ b/tests/test_nn_tensorflow_optimizations_clustering.py @@ -7,8 +7,8 @@ import math from pathlib import Path import pytest -import tensorflow as tf from flaky import flaky +from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107 from mlia.nn.tensorflow.optimizations.clustering import Clusterer from mlia.nn.tensorflow.optimizations.clustering import ClusteringConfiguration @@ -22,8 +22,8 @@ from tests.utils.common import train_model def _prune_model( - model: tf.keras.Model, target_sparsity: float, layers_to_prune: list[str] | None -) -> tf.keras.Model: + model: keras.Model, target_sparsity: float, layers_to_prune: list[str] | None +) -> keras.Model: x_train, y_train = get_dataset() batch_size = 1 num_epochs = 1 @@ -100,7 +100,7 @@ def test_cluster_simple_model_fully( """Simple MNIST test to see if clustering works correctly.""" target_sparsity = 0.5 - base_model = tf.keras.models.load_model(str(test_keras_model)) + base_model = keras.models.load_model(str(test_keras_model)) train_model(base_model) if sparsity_aware: |