wr_configs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17010eef-4033-4e37-8e1f-3ffa73056ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from wr_callback import WeightsReset, PrintEpoch\n",
    "from matplotlib import pyplot as plt\n",
    "import time\n",
    "from datasets import load_dataset, Dataset as D\n",
    "from utils import plot_history, get_dataset_name, get_csv_filename\n",
    "import csv\n",
    "\n",
    "from simple_model import make_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a4a6064-4276-4f5d-90de-090dddf6bdba",
   "metadata": {},
   "source": [
    "## Params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccfba6d4-8763-4907-8c53-a242c8871cb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = D.IMAGENETTE\n",
    "BATCH_SIZE = 32\n",
    "PENULTIMATE = 512\n",
    "EPOCHS = 80\n",
    "\n",
    "rand_configs = [\n",
    "    [1.0, 1.0, 1.0, 1.0],\n",
    "    [1.0, 1.0, 1.0, 0.5],\n",
    "    [1.0, 1.0, 0.5, 0.5],\n",
    "    [1.0, 1.0, 1.0, 0.0],\n",
    "    [1.0, 1.0, 0.5, 0.0],\n",
    "    [1.0, 0.5, 0.5, 0.0],\n",
    "    [1.0, 1.0, 0.0, 0.0],\n",
    "    [1.0, 0.5, 0.0, 0.0],\n",
    "    [0.5, 0.5, 0.0, 0.0],\n",
    "    [1.0, 0.0, 0.0, 0.0],\n",
    "    [0.5, 0.0, 0.0, 0.0],\n",
    "    [0.0, 0.0, 0.0, 0.0],\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a9a626-57b9-416e-bafa-d8caef38424e",
   "metadata": {},
   "source": [
    "## Dataset init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed868692-f363-4795-a55b-4e4d1dffbd39",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train, dataset_test, im_shape, num_classes = load_dataset(DATASET, batch_size = BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9afa8c59-f88d-491f-b9fc-389931631415",
   "metadata": {},
   "source": [
    "## Model init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80cd9212-2b53-4f5d-99e2-c990420d2b62",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = make_model(\n",
    "    im_shape, PENULTIMATE, num_classes\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825bdbd6-4e9a-47e3-b442-324ac34efb3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb1a5bf3-65ec-4426-b555-dd15dea701b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_init_weights = model.get_weights()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a158dd-d942-44f2-823c-ec294e8921c3",
   "metadata": {},
   "source": [
    "## Train model with different WR configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55a1cd24-349c-4fc4-bfe3-4755a81cdbc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "glorot_init = tf.keras.initializers.GlorotNormal()\n",
    "he_init = tf.keras.initializers.HeNormal()\n",
    "\n",
    "train_loss_values = []\n",
    "test_loss_values = []\n",
    "epochs_per_config = []\n",
    "\n",
    "training_hist_per_config = []\n",
    "\n",
    "dataset_name = get_dataset_name(DATASET)\n",
    "csv_file_name = get_csv_filename('configs', dataset_name)\n",
    "\n",
    "with open(csv_file_name, 'w+', encoding='UTF8', newline='') as f:\n",
    "    writer = csv.writer(f)\n",
    "    writer.writerow(['config', 'train loss', 'best test loss'])\n",
    "    \n",
    "for config in rand_configs:\n",
    "    print(f'---config {config}---')\n",
    "    model.set_weights(model_init_weights)\n",
    "    model.compile(\n",
    "        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),\n",
    "        loss=\"categorical_crossentropy\",\n",
    "        metrics=[\"accuracy\"],\n",
    "    )\n",
    "\n",
    "    layers_for_reset = [\n",
    "        {\"layer\": model.get_layer(name = \"dense_1\"), \"rand_lvl\": config[0], \"weights_initializer\": glorot_init},\n",
    "        {\"layer\": model.get_layer(name = \"dense\"), \"rand_lvl\": config[1], \"weights_initializer\": he_init},\n",
    "        {\"layer\": model.get_layer(name = \"conv2d_3\"), \"rand_lvl\": config[2], \"weights_initializer\": he_init},\n",
    "        {\"layer\": model.get_layer(name = \"conv2d_2\"), \"rand_lvl\": config[3], \"weights_initializer\": he_init}\n",
    "    ]\n",
    "\n",
    "    wr = WeightsReset(\n",
    "        layers_for_reset, \n",
    "        perform_reset = True, collect_stats = False, collect_weights = False, train_dataset = dataset_train)\n",
    "\n",
    "    model_hist = model.fit(\n",
    "        dataset_train,\n",
    "        epochs=EPOCHS,\n",
    "        validation_data=dataset_test,\n",
    "        callbacks=[wr,PrintEpoch()],\n",
    "        verbose=0\n",
    "    )\n",
    "\n",
    "    best_epoch = np.argmin(model_hist.history['val_loss'])\n",
    "\n",
    "    train_loss_values.append(model_hist.history['loss'][best_epoch])\n",
    "    test_loss_values.append(model_hist.history['val_loss'][best_epoch])\n",
    "    epochs_per_config.append(len(model_hist.history['val_loss']))\n",
    "    training_hist_per_config.append(model_hist.history)\n",
    "    print(f'best train loss = {train_loss_values[-1]}, best test loss = {test_loss_values[-1]}, total epochs = {epochs_per_config[-1]}')\n",
    "    print('---end---')\n",
    "    \n",
    "    with open(csv_file_name, 'a', encoding='UTF8', newline='') as f:\n",
    "        writer = csv.writer(f)\n",
    "        writer.writerow([f'{config}', train_loss_values[-1], test_loss_values[-1]])\n",
    "\n",
    "    time.sleep(30) # gpu cooler :)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41abc7e1-5423-41fc-a099-34c282adeda9",
   "metadata": {},
   "source": [
    "## Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f49e5168-9797-40c4-a8f6-b9fcba547997",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(train_loss_values)\n",
    "plt.plot(test_loss_values)\n",
    "#plt.title('model loss')\n",
    "plt.ylabel('loss')\n",
    "plt.xlabel('configuration number')\n",
    "plt.legend(['train', 'test'], loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a922c7-ed55-46f3-964e-d07f02c24231",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}