import tensorflow as tf from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Dense, InputLayer, Dropout, Conv1D, Conv2D, Flatten, Reshape, MaxPooling1D, MaxPooling2D, AveragePooling2D, BatchNormalization, Permute, ReLU, Softmax from tensorflow.keras.optimizers.legacy import Adam EPOCHS = args.epochs or 10 LEARNING_RATE = args.learning_rate or 0.0005 # If True, non-deterministic functions (e.g. shuffling batches) are not used. # This is False by default. ENSURE_DETERMINISM = args.ensure_determinism # this controls the batch size, or you can manipulate the tf.data.Dataset objects yourself BATCH_SIZE = args.batch_size or 32 if not ENSURE_DETERMINISM: train_dataset = train_dataset.shuffle(buffer_size=BATCH_SIZE*4) train_dataset=train_dataset.batch(BATCH_SIZE, drop_remainder=False) validation_dataset = validation_dataset.batch(BATCH_SIZE, drop_remainder=False) # model architecture model = Sequential() model.add(Conv2D(32, kernel_size=3, strides=2, kernel_constraint=tf.keras.constraints.MaxNorm(1), padding='same')) model.add(ReLU()) model.add(Conv2D(16, kernel_size=3, strides=2, kernel_constraint=tf.keras.constraints.MaxNorm(1), padding='same')) model.add(ReLU()) model.add(Flatten()) model.add(Dropout(0.25)) model.add(Dense(classes, name='y_pred')) model.add(Softmax()) # this controls the learning rate opt = Adam(learning_rate=LEARNING_RATE, beta_1=0.9, beta_2=0.999) callbacks.append(BatchLoggerCallback(BATCH_SIZE, train_sample_count, epochs=EPOCHS, ensure_determinism=ENSURE_DETERMINISM)) # train the neural network model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy']) model.fit(train_dataset, epochs=EPOCHS, validation_data=validation_dataset, verbose=2, callbacks=callbacks) import tensorflow as tf def akida_quantize_model( keras_model, weight_quantization: int = 4, activ_quantization: int = 4, input_weight_quantization: int = 8, ): import cnn2snn print("Performing post-training quantization...") akida_model = cnn2snn.quantize( keras_model, weight_quantization=weight_quantization, activ_quantization=activ_quantization, input_weight_quantization=input_weight_quantization, ) print("Performing post-training quantization OK") print("") return akida_model def akida_perform_qat( akida_model, train_dataset: tf.data.Dataset, validation_dataset: tf.data.Dataset, optimizer: str, fine_tune_loss: str, fine_tune_metrics: "list[str]", callbacks, stopping_metric: str = "val_accuracy", fit_verbose: int = 2, qat_epochs: int = 30, ): early_stopping = tf.keras.callbacks.EarlyStopping( monitor=stopping_metric, mode="max", verbose=1, min_delta=0, patience=10, restore_best_weights=True, ) callbacks.append(early_stopping) print("Running quantization-aware training...") akida_model.compile( optimizer=optimizer, loss=fine_tune_loss, metrics=fine_tune_metrics ) akida_model.fit( train_dataset, epochs=qat_epochs, verbose=fit_verbose, validation_data=validation_dataset, callbacks=callbacks, ) print("Running quantization-aware training OK") print("") return akida_model import tensorflow as tf def build_edge_learning_model(quantized_model, X_train, train_dataset: tf.data.Dataset, validation_dataset: tf.data.Dataset, callbacks, optimizer: str, fine_tune_loss: str, fine_tune_metrics: 'list[str]', additional_classes: int, neurons_per_class: int, num_classes: int, qat_function = None): from ei_tensorflow.brainchip.model import get_akida_converted_model import cnn2snn import akida import numpy as np from math import ceil import sys print("Looking for the feature extractor") feature_extractor_type = cnn2snn.quantization_layers.QuantizedReLU feature_extractor_found = False for layer in reversed(quantized_model.layers): if isinstance(layer, feature_extractor_type): print("") print(f"The assumed feature extractor layer is: {layer.name}") print("") print("Setting feature extractor bitwidth to 1") quantized_model = cnn2snn.quantize_layer(quantized_model, layer, bitwidth=1) feature_extractor_found = True break if not feature_extractor_found: print("EI_LOG_LEVEL=error ERROR: Can't find the feature extractor! Edge Learning model can't be built.") print("EI_LOG_LEVEL=info Try to modify 'feature_extractor_type' in the Keras Expert Mode") sys.exit(1) print("Looking for the feature extractor OK") #! After quantizing feature extractor layer to 1 bit, we need to retrain the model to recover the accuracy if qat_function: print("") print("Performing quantization-aware training...") quantized_model = qat_function(akida_model=quantized_model, train_dataset=train_dataset, validation_dataset=validation_dataset, optimizer=optimizer, fine_tune_loss=fine_tune_loss, fine_tune_metrics=fine_tune_metrics, callbacks=callbacks) print("Performing quantization-aware training OK") print("") else: print("EI_LOG_LEVEL=warn WARNING: QAT function not defined! Quantized model won't be retrained!", flush=True) akida_edge_model = get_akida_converted_model(quantized_model, MODEL_INPUT_SHAPE) #! Build edge learning compatible model akida_edge_model.pop_layer() layer_fc = akida.FullyConnected(name='akida_edge_layer', units=additional_classes * neurons_per_class, activation=False) akida_edge_model.add(layer_fc) print('Building edge compatible model OK') print('Compiling edge learning model') #! Calculate suggested number of weights as described in: #! https://doc.brainchipinc.com/examples/edge/plot_1_edge_learning_kws.html num_samples = ceil(0.1 * X_train.shape[0]) sparsities = akida.evaluate_sparsity(akida_edge_model, np.array(X_train[:num_samples]*255, dtype=np.uint8)) output_density = 1 - sparsities[akida_edge_model.layers[-2]] avg_spikes = akida_edge_model.layers[-2].output_dims[-1] * output_density #! Fix the number of weights to 1.2 times the average number of output spikes num_weights = int(1.2 * avg_spikes) print("===========================================") print(f"The number of weights is: {num_weights}") print("===========================================") try: akida_edge_model.compile(num_weights=num_weights, num_classes=additional_classes, learning_competition=0.1) except ValueError as err: print(f"EI_LOG_LEVEL=error ERROR: Can't compile Edge Learning model: {err}") print(f"EI_LOG_LEVEL=error ERROR: If the estimated number of weights is 0, try to increase the number of training cycles") sys.exit(1) print('Compiling edge learning model OK') print('') return akida_edge_model akida_model = akida_quantize_model(model) akida_model = akida_perform_qat( akida_model, train_dataset=train_dataset, validation_dataset=validation_dataset, optimizer=opt, fine_tune_loss='categorical_crossentropy', fine_tune_metrics=['accuracy'], callbacks=callbacks) akida_edge_model = build_edge_learning_model( quantized_model=akida_model, X_train=X_train, train_dataset=train_dataset, validation_dataset=validation_dataset, callbacks=callbacks, optimizer=opt, fine_tune_loss='categorical_crossentropy', fine_tune_metrics=['accuracy'], additional_classes=1, neurons_per_class=10, num_classes=classes, qat_function=akida_perform_qat)