I am creating a traditional neural network for PA behavioral modeling in Python using TensorFlow. The model takes in input I and Q values and predicts the amplifier output. So that's a .csv file with two columns. One of my goals is to prune (or optimize in any way) the model I built. The original model works just fine. However, I am facing problems while pruning the created, trained, and tested model.
Below is the original model:
import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import keras
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import Model, load_model, save_model, clone_model, Sequential
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.metrics import Accuracy
import tensorflow_model_optimization as tfmot
import time
# Load your data
data_input = pd.read_csv('Input_TimeAligned.csv',header = None)
data_input.columns = ['i','q']
data_input.describe()
data_input.head()
data_input_arr = data_input.to_numpy()
print('Input data as an array:',data_input_arr)
print(len(data_input_arr))
print()
data_output = pd.read_csv('Output_TimeAligned.csv',header = None)
data_output.columns = ['i','q']
data_output.describe()
data_output.head()
data_output_arr = data_output.to_numpy()
print('Output data as an array:',data_output_arr)
print(len(data_output_arr))
print()
data_input_tr = data_input_arr[0:122879,:]
data_output_tr = data_output_arr[0:122879,:]
data_input_test = data_input_arr[122880:491519,:]
data_output_test = data_output_arr[122880:491519,:]
X_train = data_input_tr
y_train = data_output_tr
X_test = data_input_test
y_test = data_output_test
# Define the model architecture.
start = time.time()
model = keras.Sequential([
keras.layers.InputLayer(input_shape = (2,)),
keras.layers.Dense(units = 128, activation = 'tanh', name = 'layer_1'),
keras.layers.Dense(units = 256, activation = 'tanh', name = 'layer_2'),
keras.layers.Dense(units = 512, activation = 'tanh', name = 'layer_3'),
keras.layers.Dense(units = 256, activation = 'tanh', name = 'layer_4'),
keras.layers.Dense(units = 128, activation = 'tanh', name = 'layer_5'),
keras.layers.Dense(units = 2, activation = 'tanh', name = 'output_layer'),
])
model.compile(loss = 'mean_squared_error', optimizer = 'adam', metrics = ['Accuracy'])
end = time.time()
print('Time taken to compile the model is:', end - start)
print()
print(model.summary())
print()
start = time.time()
model.fit(X_train, y_train, epochs = 3, batch_size = 32)
end = time.time()
print('Time taken to train the model is:', end - start)
print()
start = time.time()
y_hat = model.predict(X_test)
end = time.time()
print('Time taken to test the model is:', end - start)
print('The predicted output is:', y_hat)
print()
start = time.time()
model.evaluate(X_test, y_test)
end = time.time()
print('Time taken to evaluate the model is:', end - start)
print()
I tried pruning as shown below:
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
# Define model for pruning.
pruning_params = {
'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.0,
final_sparsity=0.50,
begin_step=0,
end_step=1000)
}
model_for_pruning = prune_low_magnitude(model, **pruning_params)
# `prune_low_magnitude` requires a recompile.
model_for_pruning.compile(optimizer='adam',
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model_for_pruning.summary()
start = time.time()
model.compile(
loss="mse",
optimizer=tf.keras.optimizers.Adam(learning_rate=0.001)
)
end = time.time()
print('Time taken to compile the pruned model is:', end - start)
print()
print(model.summary())
print()
start = time.time()
model.fit(
X_train,
y_train,
epochs=3,
batch_size=32,
# callbacks= pruning_callback,
# verbose=1
)
end = time.time()
print('Time taken to train the pruned model is:', end - start)
start = time.time()
y_hat = model.predict(X_test)
end = time.time()
print('Time taken to test the pruned model is:', end - start)
print('The predicted output is:', y_hat)
print()
start = time.time()
model.evaluate(X_test, y_test)
end = time.time()
print('Time taken to evaluate the pruned model is:', end - start)
print()
This is the error I am receiving:
ValueError: `prune_low_magnitude` can only prune an object of the following types: keras.models.Sequential, keras functional model, keras.layers.Layer, list of keras.layers.Layer. You passed an object of type: Sequential.
What am I doing wrong here? I am not that proficient in programming. I'd appreciate any help!
Also, am I better off doing building my model in PyTorch? Are the optimization tools there any better? Thanks again.
The error message you get is a bit misleading, prune_low_magnitude can indeed optimize Sequential models. This error happens most likely due to version incompabilities between tensorflow and tfmot. Also, installing keras on its own could resolve your issue (see here).
Furthermore, it looks a bit confusing that you are using SparseCategoricalCrossentropy loss on the pruned model. This is usually used for multi-class classification, not for regression.
Before you start "real" pruning, it might also be worth a try to reduce the number of Dense layers from 5 to 4, or reduce the number of filters in each layer.