2_size_estimator_and_profiler.ipynb

{"cells":[{"cell_type":"markdown","metadata":{"id":"Kl1A-uzhwpt_"},"source":["# **2. Model Size Estimation**\n","\n","It is no surprise that with such a tiny package, your Ardunio Nano 33 BLE Sense comes with limited memory and processing power. Therefore, you must be aware of the size and components of your model in order to have it run efficiently on your MCU.\n","\n","This notebook explores how various neural network layers affect the number of parameters, the amount memory, the number of floating point operations, and the CPU runtime of your model."]},{"cell_type":"markdown","metadata":{"id":"h5kS114ooq-0"},"source":["## 2.0 Setup GDrive and Git"]},{"cell_type":"code","execution_count":1,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":19363,"status":"ok","timestamp":1710748429902,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"},"user_tz":240},"id":"BgbZjaQZ8niT","outputId":"8a77cd0b-8158-4a79-8e21-2c651c2e70a0"},"outputs":[{"output_type":"stream","name":"stdout","text":["Mounted at /content/gdrive\n"]}],"source":["# Mount google drive\n","from google.colab import drive\n","drive.mount('/content/gdrive')"]},{"cell_type":"code","execution_count":2,"metadata":{"id":"sH-xe50YtFNd","executionInfo":{"status":"ok","timestamp":1710748431442,"user_tz":240,"elapsed":1541,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}}},"outputs":[],"source":["# Make sure your token is stored in a txt file at the location below.\n","# This way there is no risk that you will push it to your repo\n","# Never share your token with anyone, it is basically your github password!\n","with open('/content/gdrive/MyDrive/token.txt') as f:\n","    token = f.readline().strip()\n","# Use another file to store your github username\n","with open('/content/gdrive/MyDrive/git_username.txt') as f:\n","    handle = f.readline().strip()"]},{"cell_type":"code","execution_count":3,"metadata":{"id":"rRj2U4Y3ttgu","executionInfo":{"status":"ok","timestamp":1710748431442,"user_tz":240,"elapsed":3,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"adebab2c-3fb0-4927-c6c1-7f32246f51f4"},"outputs":[{"output_type":"stream","name":"stdout","text":["/content/gdrive/MyDrive/mlhs\n"]}],"source":["# Clone your github repo\n","YOUR_TOKEN = token\n","YOUR_HANDLE = handle\n","BRANCH = \"main\"\n","\n","# %mkdir /content/gdrive/MyDrive/ece5545\n","# %cd /content/gdrive/MyDrive/ece5545\n","# !git clone https://{YOUR_TOKEN}@github.com/ML-HW-SYS/a2-{YOUR_HANDLE}.git\n","# %cd /content/gdrive/MyDrive/ece5545/a2-{YOUR_HANDLE}\n","# !git checkout {BRANCH}\n","# !git pull\n","%cd /content/gdrive/MyDrive/mlhs\n","\n","PROJECT_ROOT = f\"/content/gdrive/MyDrive/mlhs/mlhs-a2\""]},{"cell_type":"code","execution_count":4,"metadata":{"id":"R9V2uP4YtzuR","executionInfo":{"status":"ok","timestamp":1710748431442,"user_tz":240,"elapsed":2,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}}},"outputs":[],"source":["# This extension reloads all imports before running each cell\n","%load_ext autoreload\n","%autoreload 2"]},{"cell_type":"markdown","metadata":{"id":"DIzgKXcDtFNe"},"source":["### Import code dependencies"]},{"cell_type":"code","execution_count":5,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":616,"status":"ok","timestamp":1710748432056,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"},"user_tz":240},"id":"ePTgb55gwT7o","outputId":"7cbe5874-1bcf-45c4-aee0-cb45f6552428"},"outputs":[{"output_type":"stream","name":"stdout","text":["/content/gdrive/MyDrive/mlhs/mlhs-a2\n","constants.py  loaders.py   __pycache__\t\tquant.py\t  train_val_test_utils.py\n","data_proc.py  networks.py  quant_conversion.py\tsize_estimate.py\n","['/content', '/env/python', '/usr/lib/python310.zip', '/usr/lib/python3.10', '/usr/lib/python3.10/lib-dynload', '', '/usr/local/lib/python3.10/dist-packages', '/usr/lib/python3/dist-packages', '/usr/local/lib/python3.10/dist-packages/IPython/extensions', '/root/.ipython']\n"]}],"source":["import sys\n","print(PROJECT_ROOT)\n","!ls {PROJECT_ROOT}/src\n","print(sys.path)"]},{"cell_type":"code","execution_count":6,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":9060,"status":"ok","timestamp":1710748441114,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"},"user_tz":240},"id":"2uKEHl-OtFNf","outputId":"5300ca2c-1778-4676-c774-0503a87a15c8"},"outputs":[{"output_type":"stream","name":"stdout","text":["Model folders are created, \n","PyTorch models will be saved in /content/gdrive/MyDrive/mlhs/models/torch_models, \n","ONNX models will be saved in /content/gdrive/MyDrive/mlhs/models/onnx_models, \n","TensorFlow Saved Models will be saved in /content/gdrive/MyDrive/mlhs/models/tf_models, \n","TensorFlow Lite models will be saved in /content/gdrive/MyDrive/mlhs/models/tflite_models, \n","TensorFlow Lite Micro models will be saved in /content/gdrive/MyDrive/mlhs/models/micro_models.\n","Imported code dependencies\n"]}],"source":["import sys,os\n","\n","# Adding assignment 2 to the system path\n","# Make sure this matches your git directory\n","sys.path.insert(0, PROJECT_ROOT)\n","\n","import torch\n","import torch.nn as nnt\n","import src.data_proc as data_proc\n","from src.constants import *\n","import numpy as np\n","\n","print(\"Imported code dependencies\")"]},{"cell_type":"markdown","metadata":{"id":"xTKG-QIfwcN9"},"source":["## 2.2 Define the Model"]},{"cell_type":"markdown","metadata":{"id":"2RRSPvUTyIQ2"},"source":["### Create the model\n","Our TinyConv model currently consists of 7 layers:\n","\n","\n","1. [Reshape](https://pytorch.org/docs/stable/generated/torch.reshape.html)\n","2. [Conv2D](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d)\n","3. [Relu](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html#torch.nn.ReLU)\n","4. [Dropout](https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout)\n","5. Reshape\n","6. [Fully Connected (Linear)](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear)\n","7. [Softmax](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html#torch.nn.Softmax)\n","\n","\n","Please refer to `<github_dir>/src/networks.py` for more detail."]},{"cell_type":"code","execution_count":7,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":844,"status":"ok","timestamp":1710748441955,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"},"user_tz":240},"id":"jU2ryBhlwcN_","outputId":"34259ccf-d990-4b8b-f34b-df3e632362f5"},"outputs":[{"output_type":"stream","name":"stdout","text":["Using cuda to run the training scrpit.\n"]}],"source":["# Define device\n","from src.networks import TinyConv\n","\n","device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","print(f'Using {device} to run the training scrpit.')"]},{"cell_type":"markdown","metadata":{"id":"FLwsdGt_liKa"},"source":["### Create data_proc.AudioProcessor() object for data preprocessing\n","When an AudioProcessor instance is created:\n","\n","1. Download speech_command dataset from DATA_URL (defined in constants.py) to data_dir (default: '/content/gdrive/MyDrive/ece5545/data')\n","default dataset url: 'https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.02.tar.gz'\n","\n","2. Determine classes and their numerical indices for training and testing based on WANTED_WORDS\n","(defined in constants.py):\n","eg. if WANTED_WORDS is ['yes', 'no'], model will be trained to identify \"yes\" and \"no\" as yes and no,\n","other words as unkown, and background noises as silence\n","\n","3. Determine and save the settings for data processing feature generator based on relavent constants\n","in constants.py\n","\n","4. Determine which audio files in the dataset are for testing, training, or validating using hash method\n","\n","5. Prepare and save background noise data using the background noise data inside dataset"]},{"cell_type":"code","execution_count":8,"metadata":{"id":"HLxcu6BB_wxg","executionInfo":{"status":"ok","timestamp":1710748919607,"user_tz":240,"elapsed":477653,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"b04f06c1-ead8-4432-868c-0ed42d73eedb"},"outputs":[{"output_type":"stream","name":"stdout","text":[">> Downloading speech_commands_v0.02.tar.gz 100.0%\n"]}],"source":["# Create audio processor (this takes some time the first time)\n","# And continues to run for a bit after reaching 100% while it's extracting files\n","audio_processor = data_proc.AudioProcessor()"]},{"cell_type":"code","execution_count":9,"metadata":{"id":"_RJ8a0otJEJp","executionInfo":{"status":"ok","timestamp":1710748919607,"user_tz":240,"elapsed":6,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"8848afbf-2b98-4d10-a62f-49bf5b7be765"},"outputs":[{"output_type":"execute_result","data":{"text/plain":["TinyConv(\n","  (conv_reshape): Reshape(output_shape=(-1, 1, 49, 40))\n","  (conv): Conv2d(1, 8, kernel_size=(10, 8), stride=(2, 2), padding=(5, 3))\n","  (relu): ReLU()\n","  (dropout): Dropout(p=0.5, inplace=False)\n","  (fc_reshape): Reshape(output_shape=(-1, 4000))\n","  (fc): Linear(in_features=4000, out_features=4, bias=True)\n","  (softmax): Softmax(dim=1)\n",")"]},"metadata":{},"execution_count":9}],"source":["# Create model\n","model_fp32 = TinyConv(audio_processor.model_settings)\n","model_fp32"]},{"cell_type":"markdown","metadata":{"id":"PcB7Ri8ewcOA"},"source":["## 2.3 Model Estimates\n","Run the next few cells to see how each layer impacts memory and runtime of the below TinyConv neural network model. Then experiment with reshaping it to see how adding or removing layers alters the metrics."]},{"cell_type":"markdown","metadata":{"id":"GEnwhmThwcOC"},"source":["### Memory Utilization\n","\n","There are two important forms of memory that we care about for MCUs: **flash memory** and **random access memory (RAM)**. Flash is **non-volatile** aka persistent storage memory; its data is saved when powered off. This is where your model's weights and code live, thus they must be able to fit within the capacity of your MCU's flash memory (1MB). On the other hand, RAM is **volatile** or non-persistent memory, thus it is used for temporary storage like input buffers and intermediate tensors. Together, they cannot exceed the size of your RAM storage (256KB).  "]},{"cell_type":"markdown","metadata":{"id":"f9Gv8LHUtFNi"},"source":["### TODO 1: Implement the `count_trainable_parameters` function in `src/size_estimate.py` to compute model size and get an estimate of the flash usage of this model\n"]},{"cell_type":"code","execution_count":10,"metadata":{"id":"--J01LrjtFNi","executionInfo":{"status":"ok","timestamp":1710748919911,"user_tz":240,"elapsed":306,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"c2dace4c-8ec6-481c-b8a6-ea34a6e934c9"},"outputs":[{"output_type":"stream","name":"stdout","text":["Total number of trainable parameters:  0.016652 M\n"]}],"source":["# Sends model weights to the GPU if tensors are on GPU\n","if torch.cuda.is_available():\n","    model_fp32.cuda()\n","\n","from src.size_estimate import count_trainable_parameters\n","num_params = count_trainable_parameters(model_fp32)\n","print(\"Total number of trainable parameters: \", num_params / float(1e6), \"M\")"]},{"cell_type":"markdown","metadata":{"id":"-ZaeZSLrtFNi"},"source":["### TODO 2: Implement the `compute_forward_size` function in `src/size_estimate.py` to compute the memory needed for a forward pass. This is how much RAM you will be using."]},{"cell_type":"code","execution_count":11,"metadata":{"id":"tLJ-LBfo56IL","executionInfo":{"status":"ok","timestamp":1710748920916,"user_tz":240,"elapsed":1006,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"dc3e505c-c4b5-4472-8711-8dfdfdcbeba8"},"outputs":[{"output_type":"stream","name":"stdout","text":["Forward memory:  0.007856 M\n"]}],"source":["# Sends model weights to the GPU if tensors are on GPU\n","if torch.cuda.is_available():\n","    model_fp32.cuda()\n","\n","from src.size_estimate import compute_forward_memory\n","frd_memory = compute_forward_memory(\n","    model_fp32,\n","    (1, model_fp32.model_settings['fingerprint_width'], model_fp32.model_settings['spectrogram_length']),\n","    device\n",")\n","print(\"Forward memory: \", frd_memory / float(1e6), \"M\")"]},{"cell_type":"markdown","metadata":{"id":"jkjm2OlXgArr"},"source":["As you can see above, the number of parameters in a neural network can add up fast which is a concern when dealing with a small amount of RAM. With the TinyConv neural network only consuming 0.21MB out of 1MB, our model will easily fit within flash memory."]},{"cell_type":"markdown","metadata":{"id":"d4qkx4L8hCBD"},"source":["### Number of Operations\n","\n","### TODO 3: Implement the `flop` function in `src/size_estimate.py` to count the total FLOPS in a forward pass with batch size = 1"]},{"cell_type":"code","execution_count":12,"metadata":{"id":"Mb5ugZA3wcOD","executionInfo":{"status":"ok","timestamp":1710748920916,"user_tz":240,"elapsed":2,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"colab":{"base_uri":"https://localhost:8080/"},"outputId":"71eae727-90e5-4bf2-e910-3ad761596ed1"},"outputs":[{"output_type":"stream","name":"stdout","text":["total number of floating operations: 672012\n","Number of FLOPs by layer and parameters:\n","Conv:  {Conv2d(1, 8, kernel_size=(10, 8), stride=(2, 2), padding=(5, 3)): 640008}\n","FC:    {Linear(in_features=4000, out_features=4, bias=True): 32004}\n"]}],"source":["from pprint import pprint\n","from src.size_estimate import flop\n","\n","if torch.cuda.is_available():\n","    model_fp32.cuda()\n","\n","# The total number of floating point operations\n","flop_by_layers = flop(\n","    model=model_fp32,\n","    input_shape=(\n","        1,\n","        model_fp32.model_settings['fingerprint_width'],\n","        model_fp32.model_settings['spectrogram_length']\n","    ),\n","    device=device)\n","total_param_flops = sum([sum(val.values()) for val in flop_by_layers.values()])\n","\n","# from thop import profile\n","# input_shape=(\n","#         1,\n","#         model_fp32.model_settings['fingerprint_width'],\n","#         model_fp32.model_settings['spectrogram_length']\n","#     )\n","# input = torch.ones(input_shape).to(device)\n","# macs, params = profile(model_fp32, inputs=(input, ))\n","# print(macs*2,params)\n","\n","print(f'total number of floating operations: {total_param_flops}')\n","print('Number of FLOPs by layer and parameters:')\n","print(\"Conv: \", flop_by_layers['conv'])\n","print(\"FC:   \", flop_by_layers['fc'])"]},{"cell_type":"markdown","metadata":{"id":"wbAbpr4Mj12i"},"source":["### CPU runtime\n","\n","### TODO 4: Measure the server/desktop CPU runtime to compare to the MCU runtime later in this assignment"]},{"cell_type":"code","execution_count":13,"metadata":{"id":"zMcAtzKHwcOD","executionInfo":{"status":"ok","timestamp":1710748920916,"user_tz":240,"elapsed":1,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}}},"outputs":[],"source":["model_fp32.cpu()\n","model_fp32.eval()\n","inputs = torch.rand([1,1960]).cpu()\n","\n","# Run a profiler to see the cpu time for inference\n","from torch.profiler import profile, record_function, ProfilerActivity\n","with profile(activities=[ProfilerActivity.CPU], record_shapes=True, with_flops=True, with_stack=True) as prof:\n","    with record_function(\"model_inference\"):\n","        model_fp32(inputs)\n","print(prof.key_averages().table(sort_by=\"cpu_time_total\", row_limit=20))"]},{"cell_type":"markdown","source":["# GPU time"],"metadata":{"id":"Qiz9Bp0cfUIi"}},{"cell_type":"code","source":["model_fp32.cuda()\n","model_fp32.eval()\n","inputs = torch.rand([1,1960]).cuda()\n","\n","# Run a profiler to see the cpu time for inference\n","from torch.profiler import profile, record_function, ProfilerActivity\n","with profile(activities=[ProfilerActivity.CUDA], record_shapes=True, with_flops=True, with_stack=True) as prof:\n","    with record_function(\"model_inference\"):\n","        model_fp32(inputs)\n","print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=20))"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"pgQRlLWTfLKy","executionInfo":{"status":"ok","timestamp":1710748921368,"user_tz":240,"elapsed":453,"user":{"displayName":"Beichen Ma","userId":"04370006166270246451"}},"outputId":"c4ff160f-8b84-4cfd-e786-f1b2d7a29107"},"execution_count":14,"outputs":[{"output_type":"stream","name":"stdout","text":["-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","void implicit_convolve_sgemm<float, float, 1024, 5, ...         0.00%       0.000us         0.00%       0.000us       0.000us      15.000us        53.57%      15.000us      15.000us             1  \n","void dot_kernel<float, 128, 0, cublasDotParams<cubla...         0.00%       0.000us         0.00%       0.000us       0.000us       4.000us        14.29%       4.000us       4.000us             1  \n","void at::native::elementwise_kernel<128, 2, at::nati...         0.00%       0.000us         0.00%       0.000us       0.000us       3.000us        10.71%       3.000us       3.000us             1  \n","void reduce_1Block_kernel<float, 128, 7, cublasGemvT...         0.00%       0.000us         0.00%       0.000us       0.000us       3.000us        10.71%       3.000us       3.000us             1  \n","void at::native::vectorized_elementwise_kernel<4, at...         0.00%       0.000us         0.00%       0.000us       0.000us       2.000us         7.14%       2.000us       2.000us             1  \n","void (anonymous namespace)::softmax_warp_forward<flo...         0.00%       0.000us         0.00%       0.000us       0.000us       1.000us         3.57%       1.000us       1.000us             1  \n","                                        cudaEventRecord         7.81%      15.000us         7.81%      15.000us      15.000us       0.000us         0.00%       0.000us       0.000us             1  \n","                                  cudaStreamIsCapturing         0.52%       1.000us         0.52%       1.000us       1.000us       0.000us         0.00%       0.000us       0.000us             1  \n","                                  cudaStreamGetPriority         0.52%       1.000us         0.52%       1.000us       1.000us       0.000us         0.00%       0.000us       0.000us             1  \n","                       cudaDeviceGetStreamPriorityRange         0.52%       1.000us         0.52%       1.000us       1.000us       0.000us         0.00%       0.000us       0.000us             1  \n","                                       cudaLaunchKernel        82.81%     159.000us        82.81%     159.000us      26.500us       0.000us         0.00%       0.000us       0.000us             6  \n","                                  cudaDeviceSynchronize         7.81%      15.000us         7.81%      15.000us      15.000us       0.000us         0.00%       0.000us       0.000us             1  \n","-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  \n","Self CPU time total: 192.000us\n","Self CUDA time total: 28.000us\n","\n"]}]}],"metadata":{"accelerator":"GPU","colab":{"gpuType":"V100","machine_shape":"hm","provenance":[{"file_id":"https://github.com/Beichen-Ma/mlhs-a2/blob/main/2_size_estimator_and_profiler.ipynb","timestamp":1710546361676}]},"interpreter":{"hash":"92bf126df007708fd70c442c808ee74575bedf7ea6317e0b182c3af0184af25d"},"kernelspec":{"display_name":"ece5545","language":"python","name":"ece5545"},"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.7.11"}},"nbformat":4,"nbformat_minor":0}