scrnet_train_models_emnist_tf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train Signal Classification and Reconstruction Network (SCRNet)\n",
    "Uses the Tensoflow machine learning library to build and train the Signal Classification and Reconstruction Network (SCRNet) to classify and reconstruct compressed Extended-MNIST <span style=\"color:blue\">***(EMNIST)*** letter</span> images with imposed Poisson noise.\n",
    "\n",
    "* First Run: **create_comp_noisy_emnist_letters_training_data.ipynb** to create compressed/noisy training data\n",
    "\n",
    "**Author:** Fabian Santiago  \n",
    "**Update:** September 18, 2024\n",
    "\n",
    "Jupyer Notebook Version: 6.5.4  \n",
    "Python Version: 3.11.5  \n",
    "TensorFlow Version: 2.16.2\n",
    "\n",
    "***[Download EMNIST](https://pypi.org/project/emnist/)***\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Modules and Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary modules and libraries \n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Input, Dense, MaxPooling2D, Conv2D, Reshape, Flatten, Add\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from emnist import extract_training_samples\n",
    "from emnist import extract_test_samples\n",
    "import numpy as np\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Values: Seed, Compression, and Output Dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the seed for TensorFlow and Keras\n",
    "tf.random.set_seed(101)\n",
    "\n",
    "# Dimension of output, original are 28 x 28\n",
    "out_dim  = 28 # Do not change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create ductionary training-sets\n",
    "in_signals_by_cmp = dict.fromkeys(['4x4','7x7','14x14','28x28'],[])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Function Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SCRNet architecture builder function\n",
    "def build_SCRNet(in_dim = 7, out_dim = 28, enc_dim = 256, learning_rate = 0.001):\n",
    "\n",
    "    # Input layer\n",
    "    inputs = Input(shape=(in_dim**2,))\n",
    "    \n",
    "    # Transform input to output signal size\n",
    "    Dec1 = Dense(out_dim**2, activation='sigmoid')(inputs)\n",
    "\n",
    "    ##################\n",
    "    # Classification #\n",
    "    ##################\n",
    "    reshpimg = Reshape((out_dim, out_dim, 1))(Dec1)\n",
    "    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(reshpimg)\n",
    "    maxp1 = MaxPooling2D((2, 2))(conv1)\n",
    "    conv2 = Conv2D(64, (3, 3), activation='relu')(maxp1)\n",
    "    maxp2 = MaxPooling2D((2, 2))(conv2)\n",
    "    flatt = Flatten()(maxp2)\n",
    "    hid1 = Dense(128, activation='relu')(flatt)\n",
    "    dpred = Dense(26, activation='softmax', name='digit-pred')(hid1)\n",
    "\n",
    "    ##########################\n",
    "    # Reconstruction Network #\n",
    "    ##########################\n",
    "    # Encoder\n",
    "    Enc2_hid = Dense(enc_dim, activation='sigmoid')(Dec1)\n",
    "    Enc2 = Dense(in_dim**2, activation='sigmoid')(Enc2_hid)\n",
    "\n",
    "    # Add dpred to match Enc2's shape\n",
    "    dpred_transformed = Dense(in_dim**2, activation='sigmoid')(dpred)\n",
    "    addenc = Add()([Enc2, dpred_transformed])\n",
    "\n",
    "    # Decoder\n",
    "    Dec2_hid = Dense(enc_dim, activation='sigmoid')(addenc)\n",
    "    reconimg = Dense(out_dim**2, activation='sigmoid', name='reconimg')(Dec2_hid)\n",
    "    \n",
    "    # Define the model\n",
    "    model = Model(inputs=inputs, outputs=[reconimg, dpred])\n",
    "\n",
    "    # Compile the model\n",
    "    # Set optimizer and learning rate\n",
    "    tf_opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
    "\n",
    "    # Compile the model\n",
    "    model.compile(optimizer=tf_opt,\n",
    "        loss={\n",
    "            'reconimg': 'mse',  # MSE for reconstruction\n",
    "            'digit-pred': 'mse'  # MSE for character recognition\n",
    "        },\n",
    "        loss_weights={\n",
    "            'reconimg': 1.00,  # Weight for the reconstruction loss\n",
    "            'digit-pred': 1.00,      # Weight for character recognition loss\n",
    "        },\n",
    "        metrics={'reconimg': 'mse', 'digit-pred': 'categorical_crossentropy'})\n",
    "    return model\n",
    "\n",
    "# Define function for fitting model using multiple stages with different epochs and batch sizes\n",
    "def fit_model(in_model, epochs_list, batch_sizes, input_data, output_data):\n",
    "\n",
    "    if len(epochs_list) != len(batch_sizes):\n",
    "        raise ValueError(f\"Input lists must be of equal size. epochs_list has length {len(epochs_list)} but batch_sizes has length {len(batch_sizes)}.\")\n",
    "    \n",
    "    for idx, epoch in enumerate(epochs_list):\n",
    "        print(f\"Training Stage #{idx+1} with {epoch} epochs and a batch size of {batch_sizes[idx]}.\\n\")\n",
    "        \n",
    "        fit_hist = in_model.fit(input_data, output_data, epochs=epoch, batch_size=batch_sizes[idx], validation_split=0.2)\n",
    "        if idx == 0:\n",
    "            fit_hist_all = fit_hist.history\n",
    "        else:\n",
    "            for key in fit_hist_all.keys():\n",
    "                fit_hist_all[key] = fit_hist_all[key] + fit_hist.history[key]\n",
    "    return in_model, fit_hist_all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Training Data\n",
    "Load traning data or load EMNIST data and create training data if it has not already been created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directory name\n",
    "directory = 'training_data'\n",
    "    \n",
    "# Load EMNIST: Used for 4x4, 7x7, 14x14, and 28x28 input architectures\n",
    "clean_train, train_labels = extract_training_samples('letters')    \n",
    "clean_test, test_labels   = extract_test_samples('letters')\n",
    "\n",
    "\n",
    "for cmp_dim in [4,7,14,28]:\n",
    "    # Set EMNIST dataset name\n",
    "    dat_file = f'{directory}/emnist_{cmp_dim}x{cmp_dim}_train.h5'\n",
    "\n",
    "    # Load compressed/noisy training data\n",
    "    with h5py.File(dat_file, 'r') as dat_file:\n",
    "        # Load compressed noisy training/test images\n",
    "        noisy_train = dat_file['noisy_train'][:]\n",
    "        noisy_test  = dat_file['noisy_test'][:]\n",
    "    \n",
    "    # Create Dictionary for compressed/noisy signals\n",
    "    in_signals = dict(zip(['train','test'],[noisy_train,noisy_test]))\n",
    "    \n",
    "    # Add to dictionary\n",
    "    in_signals_by_cmp[f'{cmp_dim}x{cmp_dim}'] = in_signals\n",
    "    \n",
    "# Prepare original mnist data for model training\n",
    "clean_train  = np.array([matrix.reshape(out_dim**2,) for matrix in clean_train/255])\n",
    "clean_test   = np.array([matrix.reshape(out_dim**2,) for matrix in clean_test/255])\n",
    "train_labels = train_labels-1\n",
    "test_labels  = test_labels-1\n",
    "train_labels = to_categorical(train_labels, num_classes=26)\n",
    "test_labels  = to_categorical(test_labels, num_classes=26)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile & Train Models: 7x7, 14x14, and 28x28\n",
    "### Train 4x4 Compressed and Gaussian Noisy Input Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create & Compile model\n",
    "scrnet4x4 = build_SCRNet(in_dim=4, out_dim=28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "# Set number of epochs (epochs_list) for each batch size (batch_sizes)\n",
    "scrnet4x4, _ = fit_model(in_model=scrnet4x4, epochs_list=[20,30], batch_sizes =[100,150], input_data = in_signals_by_cmp['4x4']['train'], output_data = [clean_train, train_labels])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train 7x7 Compressed and Noisy Input Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create & Compile model\n",
    "scrnet7x7 = build_SCRNet(in_dim=7, out_dim=28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "# Set number of epochs (epochs_list) for each batch size (batch_sizes)\n",
    "scrnet7x7, _ = fit_model(in_model=scrnet7x7, epochs_list=[20,30], batch_sizes =[100,150], input_data = in_signals_by_cmp['7x7']['train'], output_data = [clean_train, train_labels])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train 14x14 Compressed and Noisy Input Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create & Compile model\n",
    "scrnet14x14 = build_SCRNet(in_dim=14, out_dim=28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "# Set number of epochs (epochs_list) for each batch size (batch_sizes)\n",
    "scrnet14x14, _ = fit_model(in_model=scrnet14x14, epochs_list=[20,30], batch_sizes =[100,150], input_data = in_signals_by_cmp['14x14']['train'], output_data = [clean_train, train_labels])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train 28x28 Noisy Input Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create & Compile model\n",
    "scrnet28x28 = build_SCRNet(in_dim=28, out_dim=28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "# Set number of epochs (epochs_list) for each batch size (batch_sizes)\n",
    "scrnet28x28, _ = fit_model(in_model=scrnet28x28, epochs_list=[20,30], batch_sizes =[100,150], input_data = in_signals_by_cmp['28x28']['train'], output_data = [clean_train, train_labels])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save Trained Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the trained model\n",
    "scrnet4x4.save('trained_models/scrnet4x4.keras')\n",
    "scrnet7x7.save('trained_models/scrnet7x7.keras')\n",
    "scrnet14x14.save('trained_models/scrnet14x14.keras')\n",
    "scrnet28x28.save('trained_models/scrnet28x28.keras')"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}