initial commit with Panel app

Dockerfile

@@ -8,7 +8,7 @@ RUN python3 -m pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 COPY . .
 
-CMD ["panel", "serve", "/code/app.py", "--address", "0.0.0.0", "--port", "7860",  "--allow-websocket-origin", "*"]
+CMD ["panel", "serve", "/code/app.py", "--address", "0.0.0.0", "--port", "7860",  "--allow-websocket-origin", "*", "--num-procs", "2", "--num_threads", "4"]
 
 RUN mkdir /.cache
 RUN chmod 777 /.cache

app.py

@@ -1,147 +1,230 @@
-import io
-import random
-from typing import List, Tuple
-
-import aiohttp
+#####################################
+# Packages & Dependencies
+#####################################
 import panel as pn
-from PIL import Image
-from transformers import CLIPModel, CLIPProcessor
-
-pn.extension(design="bootstrap", sizing_mode="stretch_width")
-
-ICON_URLS = {
-    "brand-github": "https://github.com/holoviz/panel",
-    "brand-twitter": "https://twitter.com/Panel_Org",
-    "brand-linkedin": "https://www.linkedin.com/company/panel-org",
-    "message-circle": "https://discourse.holoviz.org/",
-    "brand-discord": "https://discord.gg/AXRHnJU6sP",
-}
-
-
-async def random_url(_):
-    pet = random.choice(["cat", "dog"])
-    api_url = f"https://api.the{pet}api.com/v1/images/search"
-    async with aiohttp.ClientSession() as session:
-        async with session.get(api_url) as resp:
-            return (await resp.json())[0]["url"]
-
-
-@pn.cache
-def load_processor_model(
-    processor_name: str, model_name: str
-) -> Tuple[CLIPProcessor, CLIPModel]:
-    processor = CLIPProcessor.from_pretrained(processor_name)
-    model = CLIPModel.from_pretrained(model_name)
-    return processor, model
-
-
-async def open_image_url(image_url: str) -> Image:
-    async with aiohttp.ClientSession() as session:
-        async with session.get(image_url) as resp:
-            return Image.open(io.BytesIO(await resp.read()))
-
-
-def get_similarity_scores(class_items: List[str], image: Image) -> List[float]:
-    processor, model = load_processor_model(
-        "openai/clip-vit-base-patch32", "openai/clip-vit-base-patch32"
-    )
-    inputs = processor(
-        text=class_items,
-        images=[image],
-        return_tensors="pt",  # pytorch tensors
-    )
-    outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image
-    class_likelihoods = logits_per_image.softmax(dim=1).detach().numpy()
-    return class_likelihoods[0]
-
-
-async def process_inputs(class_names: List[str], image_url: str):
-    """
-    High level function that takes in the user inputs and returns the
-    classification results as panel objects.
-    """
-    try:
-        main.disabled = True
-        if not image_url:
-            yield "##### ⚠️ Provide an image URL"
-            return
-    
-        yield "##### ⚙ Fetching image and running model..."
-        try:
-            pil_img = await open_image_url(image_url)
-            img = pn.pane.Image(pil_img, height=400, align="center")
-        except Exception as e:
-            yield f"##### 😔 Something went wrong, please try a different URL!"
-            return
-    
-        class_items = class_names.split(",")
-        class_likelihoods = get_similarity_scores(class_items, pil_img)
-    
-        # build the results column
-        results = pn.Column("##### 🎉 Here are the results!", img)
-    
-        for class_item, class_likelihood in zip(class_items, class_likelihoods):
-            row_label = pn.widgets.StaticText(
-                name=class_item.strip(), value=f"{class_likelihood:.2%}", align="center"
-            )
-            row_bar = pn.indicators.Progress(
-                value=int(class_likelihood * 100),
-                sizing_mode="stretch_width",
-                bar_color="secondary",
-                margin=(0, 10),
-                design=pn.theme.Material,
-            )
-            results.append(pn.Column(row_label, row_bar))
-        yield results
-    finally:
-        main.disabled = False
-
-
-# create widgets
-randomize_url = pn.widgets.Button(name="Randomize URL", align="end")
-
-image_url = pn.widgets.TextInput(
-    name="Image URL to classify",
-    value=pn.bind(random_url, randomize_url),
-)
-class_names = pn.widgets.TextInput(
-    name="Comma separated class names",
-    placeholder="Enter possible class names, e.g. cat, dog",
-    value="cat, dog, parrot",
-)
-
-input_widgets = pn.Column(
-    "##### 😊 Click randomize or paste a URL to start classifying!",
-    pn.Row(image_url, randomize_url),
-    class_names,
-)
-
-# add interactivity
-interactive_result = pn.panel(
-    pn.bind(process_inputs, image_url=image_url, class_names=class_names),
-    height=600,
-)
-
-# add footer
-footer_row = pn.Row(pn.Spacer(), align="center")
-for icon, url in ICON_URLS.items():
-    href_button = pn.widgets.Button(icon=icon, width=35, height=35)
-    href_button.js_on_click(code=f"window.open('{url}')")
-    footer_row.append(href_button)
-footer_row.append(pn.Spacer())
-
-# create dashboard
-main = pn.WidgetBox(
-    input_widgets,
-    interactive_result,
-    footer_row,
-)
-
-title = "Panel Demo - Image Classification"
-pn.template.BootstrapTemplate(
-    title=title,
-    main=main,
-    main_max_width="min(50%, 698px)",
-    header_background="#F08080",
-).servable(title=title)
+import os, yaml
+from panel.viewable import Viewer
+
+from app_components import canvas, plots
+from app_utils import styles
+
+pn.extension('plotly')
+FILE_PATH = os.path.dirname(__file__)
+
+
+################################################
+# Digit Classifier Layout
+################################################
+class DigitClassifier(Viewer):
+    '''
+    Builds and displays the UI for the classifier application. 
+
+    Args:
+        mod_path (str): The absolute path to the saved TinyVGG model
+        mod_kwargs (dict): A dictionary containing the keyword-arguments for the TinyVGG model.
+                           This should have the keys: num_blks, num_convs, in_channels, hidden_channels, and num_classes
+    '''
+
+    def __init__(self, mod_path: str, mod_kwargs: dict, **params):
+        self.canvas = canvas.Canvas(sizing_mode = 'stretch_both', 
+                                    styles = {'border':'black solid 0.15rem'})
+        
+        self.clear_btn = pn.widgets.Button(name = 'Clear', 
+                                           sizing_mode = 'stretch_width',
+                                           stylesheets = [styles.BTN_STYLESHEET])
+
+        self.plot_panels = plots.PlotPanels(canvas_info = self.canvas, mod_path = mod_path, mod_kwargs = mod_kwargs)
+        
+        super().__init__(**params)
+        self.github_logo = pn.pane.PNG(
+            object = FILE_PATH + '/assets/github-mark-white.png',
+            alt_text = 'GitHub Repo',
+            link_url = 'https://github.com/Jechen00/digit-classifier-app',
+            height = 70,
+            styles = {'margin':'0'}
+        )
+        self.controls_col = pn.FlexBox(
+            self.github_logo,
+            self.clear_btn, 
+            self.plot_panels.pred_txt,
+            gap = '60px',
+            flex_direction = 'column',
+            justify_content = 'center',
+            align_items = 'center',
+            flex_wrap = 'nowrap',
+            styles = {'width':'40%', 'height':'100%'}
+        )
+
+        self.mod_input_txt = pn.pane.HTML(
+            object = '''
+                <div>
+                    <b>MODEL INPUT</b>
+                </div>
+            ''',
+            styles = {'margin':'0rem', 'padding-left':'0.15rem', 'color':'white', 
+                      'font-size':styles.FONTSIZES['mod_input_txt'], 
+                      'font-family':styles.FONTFAMILY, 
+                      'position':'absolute', 'z-index':'100'}
+        )
+
+        self.img_row = pn.FlexBox(
+            self.canvas,
+            self.controls_col,
+            pn.FlexBox(self.mod_input_txt,
+                       self.plot_panels.img_pane, 
+                       sizing_mode = 'stretch_both',
+                       styles = {'border':'solid 0.15rem white'}),
+            gap = '1%',
+            flex_wrap = 'nowrap',
+            flex_direction = 'row',
+            justify_content = 'center',
+            sizing_mode = 'stretch_width',
+            styles = {'height':'60%'}
+        )
+
+        self.prob_row =  pn.FlexBox(self.plot_panels.prob_pane,
+                                    sizing_mode = 'stretch_width',
+                                    styles = {'height':'40%', 
+                                              'border':'solid 0.15rem black'})
+
+        self.page_info = pn.pane.HTML(
+            object = f'''
+                <style>
+                    .link {{
+                        color: rgb(29, 161, 242);
+                        text-decoration: none;
+                        transition: text-decoration 0.2s ease;
+                    }}
+
+                    .link:hover {{
+                        text-decoration: underline;
+                    }}
+                </style>
+
+                <div style="text-align:center; font-size:{styles.FONTSIZES['sidebar_title']};margin-top:0.2rem">
+                    <b>Digit Classifier</b>
+                </div>
+
+                <div style="padding:0 2.5% 0 2.5%; text-align:left; font-size:{styles.FONTSIZES['sidebar_txt']}; width: 100%;">
+                    <hr style="height:2px; background-color:rgb(200, 200, 200); border:none; margin-top:0">
+
+                    <p style="margin:0">
+                        This is a handwritten digit classifier that uses a <i>convolutional neural network (CNN)</i>
+                        to make predictions. The architecture of the model is a scaled-down version of 
+                        the <i>Visual Geometry Group (VGG)</i> architecture from the paper:
+                        <a href="https://arxiv.org/pdf/1409.1556" 
+                           class="link" 
+                           target="_blank" 
+                           rel="noopener noreferrer">
+                        Very Deep Convolutional Networks for Large-Scale Image Recognition</a>.
+                    </p>
+                    </br>
+                    <p style="margin:0">
+                        <b>How To Use:</b> Draw a digit (0-9) on the canvas 
+                        and the model will produce a prediction for it in real time.
+                        Prediction probabilities (or confidences) for each digit are displayed in the bar chart, 
+                        reflecting the model's softmax output distribution.
+                        To the right of the canvas, you'll also find the transformed input image, i.e. the canvas drawing after undergoing  
+                        <a href="https://paperswithcode.com/dataset/mnist"
+                           class="link" 
+                           target="_blank" 
+                           rel="noopener noreferrer">
+                        MNIST preprocessing</a>.  
+                        This input image represents what the model receives prior to feature extraction and classification.
+                    </p>
+                </div>
+                <div style="margin-left: 5px; margin-top: 72px">
+                    <a href="https://github.com/Jechen00"
+                    class="link"
+                    target="blank"
+                    rel="noopener noreferrer"
+                    style="font-size: {styles.FONTSIZES['made_by_txt']}; color: {styles.CLRS['made_by_txt']};">
+                        Made by Jeff Chen
+                    </a>
+                </div>
+            ''',
+            styles = {'margin':' 0rem', 'color': styles.CLRS['sidebar_txt'], 
+                      'width': '19.7%', 'height': '100%',
+                      'font-family': styles.FONTFAMILY,
+                      'background-color': styles.CLRS['sidebar'],
+                      'overflow-y':'scroll',
+                      'border': 'solid 0.15rem black'}
+        )
+
+        self.classifier_content = pn.FlexBox(
+            self.img_row,
+            self.prob_row,
+            gap = '0.5%',
+            flex_direction = 'column',
+            flex_wrap = 'nowrap',
+            sizing_mode = 'stretch_height',
+            styles = {'width': '80%'}
+        )
+
+        self.page_content = pn.FlexBox(
+            self.page_info,
+            self.classifier_content,
+            gap = '0.3%',
+            flex_direction = 'row',
+            justify_content = 'space-around',
+            align_items = 'center',
+            flex_wrap = 'nowrap',
+            styles = {
+                'height':'100%',
+                'width':'100vw',
+                'padding': '1%',
+                'min-width': '1200px',
+                'min-height': '600px',
+                'max-width': '3600px',
+                'max-height': '1800px',
+                'background-color': styles.CLRS['page_bg']
+            },
+        )
+
+        # This is mainly used to ensure there is always have a grey background
+        self.page_layout = pn.FlexBox(
+            self.page_content,
+            justify_content = 'center',
+            flex_wrap = 'nowrap',
+            sizing_mode = 'stretch_both',
+            styles = {
+                'min-width': 'max-content',
+                'background-color': styles.CLRS['page_bg'],
+            }
+        )
+        # Set up on-click event with clear button and the canvas
+        self.clear_btn.on_click(self.canvas.toggle_clear)
+        
+    def __panel__(self):
+        '''
+        Returns the main layout of the application to be rendered by Panel.
+        '''
+        return self.page_layout
+
+
+def create_app():
+    '''
+    Creates the application, ensuring that each user gets a different instance of digit_classifier.
+    Mostly used to keep things away from a global scope.
+    '''
+    # Used to serve with panel serve in command line
+    save_dir = FILE_PATH + '/saved_models'
+    base_name = 'tiny_vgg_less_compute'
+
+    mod_path = f'{save_dir}/{base_name}_model.pth' # Path to the saved model state dict
+    settings_path = f'{save_dir}/{base_name}_settings.yaml' # Path to the saved model kwargs
+
+    # Load in model kwargs
+    with open( settings_path, 'r') as f:
+        loaded_settings = yaml.load(f, Loader = yaml.FullLoader)
+
+    mod_kwargs = loaded_settings['mod_kwargs']
+
+    digit_classifier = DigitClassifier(mod_path = mod_path, mod_kwargs = mod_kwargs)
+    return digit_classifier
+
+################################################
+# Serve App
+################################################
+# Used to serve with panel serve in command line
+create_app().servable(title = 'CNN Digit Classifier')

app_components/canvas.py

@@ -0,0 +1,84 @@
+#####################################
+# Packages & Dependencies
+#####################################
+import param
+from panel.reactive import ReactiveHTML
+
+
+#####################################
+# Canvas
+#####################################
+class Canvas(ReactiveHTML):
+    '''
+    The HTML canvas panel used for drawing digits (0-9) in the application.
+    Reference: https://panel.holoviz.org/how_to/custom_components/examples/canvas_draw.html
+    '''
+    uri = param.String()
+    clear = param.Boolean(default = False)
+    
+    _template = '''
+        <canvas
+          id="canvas"
+          style="width: 100%; height: 100%"
+          height=400px
+          width=400px
+          onmousedown="${script('start')}"
+          onmousemove="${script('draw')}"
+          onmouseup="${script('end')}"
+          onmouseleave="${script('end')}">
+        </canvas>
+    '''
+    
+    _scripts = {
+        'render': '''
+            state.ctx = canvas.getContext('2d');
+            state.ctx.fillStyle = '#FFFFFF';
+            state.ctx.fillRect(0, 0, canvas.width, canvas.height);
+            state.ctx.lineWidth = 30;
+            state.ctx.strokeStyle = '#000000';
+            state.ctx.lineJoin = 'round';
+            state.ctx.lineCap = 'round';
+
+            // Helper to normalize mouse coordinates
+            state.getCoords = function(e) {
+                const rect = canvas.getBoundingClientRect();
+                return {
+                    x: (e.clientX - rect.left) * (canvas.width / rect.width),
+                    y: (e.clientY - rect.top) * (canvas.height / rect.height)
+                };
+            };
+        ''',
+        
+       'start': '''
+            if (state.isDrawing) return;
+            state.isDrawing = true;
+            const pos = state.getCoords(event);
+            state.ctx.beginPath();
+            state.ctx.moveTo(pos.x, pos.y);
+        ''',
+        
+        'draw': '''
+            if (!state.isDrawing) return;
+            const pos = state.getCoords(event);
+            state.ctx.lineTo(pos.x, pos.y);
+            state.ctx.stroke();
+            data.uri = canvas.toDataURL('image/png');
+        ''',
+        
+        'end': '''
+            if (!state.isDrawing) return;      // Early return if already not drawing
+            state.isDrawing = false;
+        ''',
+        
+        'clear': '''
+            state.ctx.fillStyle = '#FFFFFF';
+            state.ctx.fillRect(0, 0, canvas.width, canvas.height);
+            data.uri = '';
+        '''
+    }
+
+    def toggle_clear(self, *event):
+        '''
+        Toggles the value of self.clear to trigger the JS 'clear' function.
+        '''
+        self.clear = not self.clear

app_components/plots.py

@@ -0,0 +1,225 @@
+#####################################
+# Packages & Dependencies
+#####################################
+import param
+import panel as pn
+
+import torch
+import numpy as np
+import plotly.graph_objects as go
+
+from . import canvas
+from app_utils import styles
+
+import sys, os
+APP_PATH = os.path.dirname(os.path.dirname(__file__)) # Path to the digit-classifier-app directory
+sys.path.append(APP_PATH + '/model_training')
+
+# Imports from model_training
+import data_setup, model
+
+
+#####################################
+# Plotly Panels
+#####################################
+PLOTLY_CONFIGS = {
+    'displayModeBar': True, 'displaylogo': False,
+    'modeBarButtonsToRemove': ['autoScale', 'lasso', 'select', 
+                               'toImage', 'pan', 'zoom', 'zoomIn', 'zoomOut']
+}
+
+class PlotPanels(param.Parameterized):
+    '''
+    Contains all Plotly pane objects for the application. 
+    This includes the probability bar chart and the MNIST preprocessed image heat map.
+
+    Args:
+        canvas_info (param.ClassSelector): A Canvas class object to get the data URI of the drawn image.
+        mod_path (str): The absolute path to the saved TinyVGG model.
+        mod_kwargs (dict): A dictionary containing the keyword-arguments for the TinyVGG model.
+                           This should have the keys: num_blks, num_convs, in_channels, hidden_channels, and num_classes
+    '''
+
+    canvas_info = param.ClassSelector(class_ = canvas.Canvas)    # Canvas object to get the data URI 
+    
+    def __init__(self, mod_path: str, mod_kwargs: dict, **params):
+        super().__init__(**params)
+        self.class_labels = np.arange(0, 10)
+        self.cnn_mod = model.TinyVGG(**mod_kwargs)
+        self.cnn_mod.load_state_dict(torch.load(mod_path, map_location = 'cpu'))
+        
+        self.img_pane = pn.pane.Plotly(
+            name = 'image_plot',
+            config = PLOTLY_CONFIGS,
+            sizing_mode = 'stretch_both',
+            margin = 0,
+        )
+
+        self.prob_pane = pn.pane.Plotly(
+            name = 'prob_plot',
+            config = PLOTLY_CONFIGS,
+            sizing_mode = 'stretch_both',
+            margin = 0
+        )
+        
+        self.pred_txt = pn.pane.HTML(
+            styles = {'margin':'0rem', 'color':styles.CLRS['pred_txt'], 
+                      'font-size':styles.FONTSIZES['pred_txt'],
+                      'font-family':styles.FONTFAMILY}
+        )
+
+        # Initialize plotly figures
+        self._update_prediction()
+
+        # Set up watchers thta update based on data URI changes
+        self.canvas_info.param.watch(self._update_prediction, 'uri')
+
+    def _update_prediction(self, *event):
+        '''
+        Performs all prediction-related updates for the application.
+        This function is connected to the URI parameter of canvas_info through a watcher.
+        Any times the URI changes, a class prediction is immediately. 
+        Following this, the probability bar chart and model input heatmap are updated as well.
+        '''
+        self._update_preprocessed_tensor()
+        self._update_pred_txt()
+        self._update_img_plot()
+        self._update_prob_plot()
+
+    def _update_preprocessed_tensor(self):
+        '''
+        Transforms the data URI (string) from canvas_info into a preprocessed tensor.
+        This is done by having it undergo the MNISt preprocessing pipeline (see mnist_preprocess in data_setup for details).
+        Additionally, a prediction is made for the preprocessed tensor to get its class label. 
+        The correpsonding set of prediction probabilities are stored.
+        '''
+        # Check if uri is non-empty
+        if self.canvas_info.uri:
+            self.input_img = data_setup.mnist_preprocess(self.canvas_info.uri)
+
+            self.cnn_mod.eval() # Set CNN to eval & inference mode
+            with torch.inference_mode():
+                pred_logits = self.cnn_mod(self.input_img.unsqueeze(0))
+                self.pred_probs = torch.softmax(pred_logits, dim = 1)[0].numpy()
+                self.pred_label = np.argmax(self.pred_probs)
+        else:
+            self.input_img = torch.zeros((28, 28))
+            self.pred_probs = np.zeros(10)
+            self.pred_label = None
+
+    def _update_pred_txt(self):
+        '''
+        Updates the prediction and probability HTML text to reflect the current data URI.
+        '''
+        if self.canvas_info.uri:
+            pred, prob = self.pred_label, f'{self.pred_probs[self.pred_label]:.3f}'
+        else:
+            pred, prob = 'N/A', 'N/A'
+
+        self.pred_txt.object = f'''
+            <div style="text-align: left;">
+                <b>Prediction:</b> {pred}
+                </br>
+                <b>Probability:</b> {prob}
+            </div>
+        '''
+
+    def _update_prob_plot(self):
+        '''
+        Updates the probability bar chart to showcase the softmax output probability distribution
+        obtained from the prediction in _update_preprocessed_tensor.
+        '''
+        # Marker fill and outline color for bar plot
+        mkr_clrs = [styles.CLRS['base_bar']] * len(self.class_labels)
+        mkr_line_clrs = [styles.CLRS['base_bar_line']] * len(self.class_labels)
+        if self.pred_label is not None:
+            mkr_clrs[self.pred_label] = styles.CLRS['pred_bar']
+            mkr_line_clrs[self.pred_label] = styles.CLRS['pred_bar_line']
+            
+        fig = go.Figure()
+        # Bar plot
+        fig.add_trace(
+            go.Bar(x = self.class_labels, y = self.pred_probs, 
+                   marker_color = mkr_clrs, marker_line_color = mkr_line_clrs,
+                   marker_line_width = 1.5, showlegend = False,
+                   text = self.pred_probs, textposition = 'outside',
+                   textfont = dict(color = styles.CLRS['plot_txt'],
+                                   size = styles.FONTSIZES['plot_bar_txt'], family = styles.FONTFAMILY), 
+                   texttemplate = '%{text:.3f}', 
+                   customdata = self.pred_probs * 100,
+                   hoverlabel_font = dict(family = styles.FONTFAMILY),
+                   hovertemplate = '<b>Class Label:</b> %{x}' +
+                                   '<br><b>Probability:</b> %{customdata:.2f} %' +
+                                   '<extra></extra>'
+            )
+        )
+        # Used to fix axis limits
+        fig.add_trace(
+            go.Scatter(
+                x = [0.5, 0.5], y = [0.1, 1],
+                marker = dict(color = 'rgba(0, 0, 0, 0)', size = 10),
+                mode = 'markers', 
+                hoverinfo = 'skip', 
+                showlegend = False
+            )
+        )
+        fig.update_yaxes(
+            title = dict(text = 'Prediction Probability', standoff = 0,
+                         font = dict(color = styles.CLRS['plot_txt'],
+                                     size = styles.FONTSIZES['plot_labels'], 
+                                     family = styles.FONTFAMILY)),
+            tickfont = dict(size = styles.FONTSIZES['plot_ticks'], 
+                            family = styles.FONTFAMILY),
+            dtick = 0.1, ticks = 'outside', ticklen = 0,
+            gridcolor = styles.CLRS['prob_plot_grid']
+        )
+        fig.update_xaxes(
+            title = dict(text = 'Class Label', standoff = 6,
+                         font = dict(color = styles.CLRS['plot_txt'],
+                                     size = styles.FONTSIZES['plot_labels'], 
+                                     family = styles.FONTFAMILY)),
+            dtick = 1, tickfont = dict(size = styles.FONTSIZES['plot_ticks'], 
+                                       family = styles.FONTFAMILY),
+        )
+        fig.update_layout(
+            paper_bgcolor = styles.CLRS['prob_plot_bg'],
+            plot_bgcolor = styles.CLRS['prob_plot_bg'],
+            margin = dict(l = 60, r = 0, t = 5, b = 45),
+        )
+
+        self.prob_pane.object = fig
+        
+    def _update_img_plot(self):
+        '''
+        Updates the heat map to showcase the current model input, i.e. the preprocessed canvas drawing.
+        '''
+        img_np = self.input_img.squeeze().numpy()
+
+        if self.pred_label is not None:
+            zmin, zmax = np.min(img_np), np.max(img_np)
+        else:
+            zmin, zmax = 0, 1
+
+        fig = go.Figure(
+            data = go.Heatmap(
+                z = img_np,
+                colorscale = 'gray',
+                showscale = False,
+                zmin = zmin,
+                zmax = zmax,
+                hoverlabel_font = dict(family = styles.FONTFAMILY),
+                hovertemplate = '<b>Pixel Position:</b> (%{x}, %{y})' +
+                                '<br><b>Pixel Value:</b> %{z:.3f}' + 
+                                '<extra></extra>'
+            )
+        )
+
+        fig.update_yaxes(autorange = 'reversed') 
+        fig.update_layout(
+            plot_bgcolor = styles.CLRS['img_plot_bg'],
+            margin = dict(l = 0, r = 0, t = 0, b = 0),
+            xaxis = dict(showticklabels = False),
+            yaxis = dict(showticklabels = False),
+        )
+
+        self.img_pane.object = fig

app_utils/styles.py

@@ -0,0 +1,53 @@
+#####################################
+# Fonts & Colors
+#####################################
+FONTFAMILY = 'Helvetica'
+
+FONTSIZES = {
+    'pred_txt': '1.2rem',
+    'mod_input_txt': '0.8rem',
+    'plot_ticks': 14,
+    'plot_labels': 16,
+    'plot_bar_txt': 14,
+    'btn': '1rem',
+    'sidebar_txt': '0.95rem',
+    'sidebar_title': '1.8rem',
+    'made_by_txt': '0.75rem'
+}
+
+CLRS = {
+    'pred_txt': 'white',
+    'sidebar': 'white',
+    'sidebar_txt': 'black',
+    'base_bar': 'rgb(158, 202, 225)',
+    'base_bar_line': 'rgb(8, 48, 107)',
+    'pred_bar': 'rgb(240, 140, 140)',
+    'pred_bar_line': 'rgb(180, 0, 0)',
+    'plot_txt': 'black',
+    'prob_plot_bg': 'white',
+    'prob_plot_grid': 'rgb(225, 225, 225)',
+    'img_plot_bg': 'black',
+    'btn_base': 'white',
+    'btn_hover': 'rgb(200, 200, 200)',
+    'page_bg': 'rgb(150, 150, 150)',
+    'made_by_txt': 'rgb(180, 180, 180)'
+}
+
+
+#####################################
+# Stylesheets
+#####################################
+BTN_STYLESHEET = f'''
+    :host(.solid) .bk-btn {{
+        background-color: {CLRS['btn_base']};
+        border: black solid 0.1rem;
+        border-radius: 0.8rem;
+        font-size: {FONTSIZES['btn']};
+        padding-top: 0.3rem;
+        padding-bottom: 0.3rem;
+    }} 
+    
+    :host(.solid) .bk-btn:hover {{
+        background-color: {CLRS['btn_hover']};
+    }}
+'''

model_training/args.txt

@@ -0,0 +1,12 @@
+--num-workers
+0
+--num-epochs
+25
+--batch-size
+100
+--learning-rate
+0.001
+--patience
+10
+--min-delta
+0.001

model_training/data_setup.py

@@ -0,0 +1,151 @@
+#####################################
+# Packages & Dependencies
+#####################################
+from torchvision import transforms, datasets
+from torch.utils.data import DataLoader
+
+import utils
+from typing import Tuple
+
+import io
+import base64
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Transformations applied to each image
+BASE_TRANSFORMS = transforms.Compose([
+    transforms.ToTensor(), # Convert to tensor and rescale pixel values to within [0, 1]
+    transforms.Normalize(mean = [0.1307], std = [0.3081]) # Normalize with MNIST stats
+])
+
+TRAIN_TRANSFORMS = transforms.Compose([
+    transforms.RandomAffine(degrees = 15, # Rotate up to -/+ 15 degrees
+                            scale = (0.8, 1.2), # Scale between 80 and 120 percent
+                            translate = (0.08, 0.08), # Translate up to -/+ 8 percent in both x and y
+                            shear = 10),  # Shear up to -/+ 10 degrees
+    transforms.ToTensor(), # Convert to tensor and rescale pixel values to within [0, 1]
+    transforms.Normalize(mean = [0.1307], std = [0.3081]), # Normalize with MNIST stats
+])
+
+
+#####################################
+# Functions
+#####################################
+def get_dataloaders(root: str, 
+                    batch_size: int, 
+                    num_workers: int = 0) -> Tuple[DataLoader, DataLoader]:
+    '''
+    Creates training and testing dataloaders for the MNIST dataset
+
+    Args:
+        root (str): Path to download MNIST data.
+        batch_size (int): Size used to split training and testing datasets into batches.
+        num_workers (int): Number of workers to use for multiprocessing. Default is 0.
+    '''
+
+    # Get training and testing MNIST data
+    mnist_train = datasets.MNIST(root, download = True, train = True, 
+                                transform = TRAIN_TRANSFORMS)
+    mnist_test = datasets.MNIST(root, download = True, train = False, 
+                                transform = BASE_TRANSFORMS)
+
+    # Create dataloaders
+    if num_workers > 0:
+        mp_context = utils.MP_CONTEXT
+    else:
+        mp_context = None
+
+    train_dl = DataLoader(
+        dataset = mnist_train,
+        batch_size = batch_size,
+        shuffle = True,
+        num_workers = num_workers,
+        multiprocessing_context = mp_context,
+        pin_memory = True
+    )
+
+    test_dl = DataLoader(
+        dataset = mnist_test,
+        batch_size = batch_size,
+        shuffle = False,
+        num_workers = num_workers,
+        multiprocessing_context = mp_context,
+        pin_memory = True
+    )
+
+    return train_dl, test_dl
+
+def mnist_preprocess(uri: str, plot: bool = False):
+    '''
+    Preprocesses a data URI representing a handwritten digit image according to the pipeline used in the MNIST dataset.
+    The pipeline includes:
+        1. Converting the image to grayscale.
+        2. Resizing the image to 20x20, preserving the aspect ratio, and using anti-aliasing.
+        3. Centering the resized image in a 28x28 image based on the center of mass (COM).
+        4. Converting the image to a tensor (pixel values between 0 and 1) and normalizing it using MNIST statistics.
+
+    Reference: https://paperswithcode.com/dataset/mnist
+
+    Args:
+        uri (str): A string representing the full data URI.
+        plot (bool, optional): If True, the resized 20x20 image is plotted alongside the final 28x28 image (pre-normalization). 
+                               The red lines on these plots intersect at the COM position. Default is False.
+    Returns:
+        Tensor: A tensor of shape (1, 28, 28) representing the preprocessed image, normalized using MNIST statistics.
+    '''
+    encoded_img = uri.split(',', 1)[1]
+    image_bytes = io.BytesIO(base64.b64decode(encoded_img))
+    pil_img = Image.open(image_bytes).convert('L') # Gray scale
+    
+    # Resize to 20x20, preserving aspect ratio, and using anti-aliasing
+    pil_img.thumbnail((20, 20), Image.Resampling.LANCZOS)
+
+    # Convert to numpy and invert image
+    img = 255 - np.array(pil_img)
+
+    # Get image indices for y-axis (rows) and x-axis (columns)
+    img_idxs = np.indices(img.shape)
+    tot_mass = img.sum()
+    
+    # This represents the indices of the center of masses (COMs)
+    com_x = np.round((img_idxs[1] * img).sum() / tot_mass).astype(int)
+    com_y = np.round((img_idxs[0] * img).sum() / tot_mass).astype(int)
+    
+    dist_com_end_x = img.shape[1] - com_x # number of column indices from com_x to last index
+    dist_com_end_y = img.shape[0] - com_y # number of row indices from com_y to last index
+    
+    new_img = np.zeros((28, 28), dtype = np.uint8)
+    new_com_x, new_com_y = 14, 14 # Indices of the COMs for the new 28x28 image
+    
+    valid_start_x = min(new_com_x, com_x)
+    valid_end_x = min(14, dist_com_end_x) # 14 is index distance from new COM to 28-th index
+    valid_start_y = min(new_com_y, com_y)
+    valid_end_y = min(14, dist_com_end_y) # 14 is index distance from new COM to 28-th index
+    
+    old_slice_x = slice(com_x - valid_start_x, com_x + valid_end_x)
+    old_slice_y = slice(com_y - valid_start_y, com_y + valid_end_y)
+    new_slice_x = slice(new_com_x - valid_start_x, new_com_x + valid_end_x)
+    new_slice_y = slice(new_com_y - valid_start_y, new_com_y + valid_end_y)
+
+    # Paste cropped image into 28x28 field such that the old COM (com_y, com_x), is at the center (14, 14)
+    new_img[new_slice_y, new_slice_x] = img[old_slice_y, old_slice_x]
+    
+    if plot:
+        fig, axes = plt.subplots(nrows = 1, ncols = 2, figsize = (12, 6))
+
+        axes[0].imshow(img, cmap = 'grey')
+        axes[0].axhline(com_y, c = 'red')
+        axes[0].axvline(com_x, c = 'red')
+
+        axes[1].imshow(new_img, cmap = 'grey')
+        axes[1].axhline(new_com_y, c = 'red')
+        axes[1].axvline(new_com_x, c = 'red')
+        
+        axes[0].set_title(f'Original Resized {img.shape[0]}x{img.shape[1]} Image')
+        axes[1].set_title('New Centered 28x28 Image')
+        
+        plt.tight_layout()
+
+    # Return transformed tensor of new image. This includes normalizing to MNIST stats
+    return BASE_TRANSFORMS(new_img)

model_training/engine.py

@@ -0,0 +1,195 @@
+#####################################
+# Packages
+#####################################
+import torch
+
+from typing import Tuple, Dict, List
+import utils
+
+#####################################
+# Functions
+#####################################
+def train_step(model: torch.nn.Module,
+               dataloader: torch.utils.data.DataLoader,
+               loss_fn: torch.nn.Module,
+               optimizer: torch.optim.Optimizer,
+               device: torch.device) -> Tuple[float, float]:
+    
+    '''
+    Performs a training step for a PyTorch model.
+    
+    Args:
+        model (torch.nn.Module): PyTorch model that will be trained
+        dataloader (torch.utils.data.DataLoader): Dataloader containing data to train on
+        loss_fn (torch.nn.Module): Loss function used as the error metric
+        optimizer (torch.optim.Optimizer): Optimization method used to update model parameters per batch
+        device (torch.device): Device to train on
+        
+    Returns:
+        train_loss (float): The average loss calculated over the training set.
+        train_acc (float): The accuracy calculated over the training set.
+    '''
+    
+    model.train()
+    train_loss = torch.tensor(0.0, device = device)
+    train_acc = torch.tensor(0.0, device = device)
+    num_samps = len(dataloader.dataset)
+
+    # Loop through all batches in the dataloader
+    for X, y in dataloader:
+        
+        optimizer.zero_grad() # Clear old accumulated gradients
+        
+        X, y = X.to(device), y.to(device)
+        
+        y_logits = model(X) # Get logits
+        
+        loss = loss_fn(y_logits, y)
+        train_loss += loss.detach() * X.shape[0] # Calculate total loss for batch
+        
+        loss.backward() # Perform backpropagation
+        optimizer.step() # Update parameters
+        
+        y_pred = y_logits.argmax(dim = 1) # No softmax needed for argmax (b/c preserves order)
+        
+        train_acc += (y_pred == y).sum() # Calculate total accuracy for batch
+    
+    # Get average loss and accuracy per sample
+    train_loss = train_loss.item() / num_samps
+    train_acc = train_acc.item() / num_samps
+    
+    return train_loss, train_acc
+
+
+def test_step(model: torch.nn.Module,
+              dataloader: torch.utils.data.DataLoader,
+              loss_fn: torch.nn.Module,
+              device: torch.device) -> Tuple[float, float]:
+    
+    '''
+    Performs a testing step for a PyTorch model.
+    
+    Args:
+        model (torch.nn.Module): PyTorch model that will be tested.
+        dataloader (torch.utils.data.DataLoader): Dataloader containing data to test on.
+        loss_fn (torch.nn.Module): Loss function used as the error metric.
+        device (torch.device): Device to compute on.
+        
+    Returns:
+        test_loss (float): The average loss calculated over batches.
+        test_acc (float): The average accuracy calculated over batches.
+    '''
+    
+    model.eval()
+    test_loss = torch.tensor(0.0, device = device)
+    test_acc = torch.tensor(0.0, device = device)
+    num_samps = len(dataloader.dataset)
+
+    with torch.inference_mode():
+        # Loop through all batches in the dataloader
+        for X, y in dataloader:
+            X, y = X.to(device), y.to(device)
+
+            y_logits = model(X) # Get logits
+
+            test_loss += loss_fn(y_logits, y) * X.shape[0] # Calculate total loss for batch
+
+            y_pred = y_logits.argmax(dim = 1) # No softmax needed for argmax (b/c preserves order)
+
+            test_acc += (y_pred == y).sum() # Calculate total accuracy for batch
+
+    # Get average loss and accuracy
+    test_loss = test_loss.item() / num_samps
+    test_acc = test_acc.item() / num_samps
+    
+    return test_loss, test_acc
+
+
+def train(model: torch.nn.Module,
+          train_dl: torch.utils.data.DataLoader,
+          test_dl: torch.utils.data.DataLoader,
+          loss_fn: torch.nn.Module,
+          optimizer: torch.optim.Optimizer,
+          num_epochs: int,
+          patience: int,
+          min_delta: float,
+          device: torch.device,
+          save_mod: bool = True,
+          save_dir: str = '',
+          mod_name: str = '') -> Dict[str, List[float]]:
+    '''
+    Performs the training and testing steps for a PyTorch model, 
+    with early stopping applied for test loss.
+    
+    Args:
+        model (torch.nn.Module): PyTorch model to train.
+        train_dl (torch.utils.data.DataLoader): DataLoader for training.
+        test_dl (torch.utils.data.DataLoader): DataLoader for testing.
+        loss_fn (torch.nn.Module): Loss function used as the error metric.
+        optimizer (torch.optim.Optimizer): Optimizer used to update model parameters per batch.
+        
+        num_epochs (int): Max number of epochs to train.
+        patience (int): Number of epochs to wait before early stopping.
+        min_delta (float): Minimum decrease in loss to reset counter.
+        
+        device (torch.device): Device to train on.
+        save_mod (bool, optional): If True, saves the model after each epoch. Default is True.
+        save_dir (str, optional): Directory to save the model to. Must be nonempty if save_mod is True.
+        mod_name (str, optional): Filename for the saved model. Must be nonempty if save_mod is True.
+
+    returns:
+        res (dict): A results dictionary containing lists of train and test metrics for each epoch.
+    '''
+    
+    bold_start, bold_end = '\033[1m', '\033[0m'
+
+    if save_mod:
+        assert save_dir, 'save_dir cannot be None or empty.'
+        assert mod_name, 'mod_name cannot be None or empty.'
+
+    # Initialize results dictionary
+    res = {'train_loss': [],
+           'train_acc': [],
+           'test_loss': [],
+           'test_acc': []
+    }
+    
+    # Initialize best_loss and counter for early stopping
+    best_loss, counter = None, 0
+        
+    for epoch in range(num_epochs):
+        # Perform training and testing step
+        train_loss, train_acc = train_step(model, train_dl, loss_fn, optimizer, device)
+        test_loss, test_acc = test_step(model, test_dl, loss_fn, device)
+        
+        # Store loss and accuracy values
+        res['train_loss'].append(train_loss)
+        res['train_acc'].append(train_acc)
+        res['test_loss'].append(test_loss)
+        res['test_acc'].append(test_acc)
+        
+        print(f'Epoch: {epoch + 1} | ' +
+             f'train_loss = {train_loss:.4f} | train_acc = {train_acc:.4f} | ' +
+             f'test_loss = {test_loss:.4f} | test_acc = {test_acc:.4f}')
+        
+        # Check for improvement
+        if best_loss == None:
+            best_loss = test_loss
+            if save_mod:
+                utils.save_model(model, save_dir, mod_name)
+
+        elif test_loss < best_loss - min_delta:
+            best_loss = test_loss
+            counter = 0
+
+            if save_mod:
+                utils.save_model(model, save_dir, mod_name)
+                print(f'{bold_start}[SAVED]{bold_end} Adequate improvement in test loss; model saved.')
+
+        else:
+            counter += 1
+            if counter > patience:
+                print(f'{bold_start}[ALERT]{bold_end} No improvement in test loss after {counter} epochs; early stopping triggered.')
+                break
+
+    return res

model_training/model.py

@@ -0,0 +1,122 @@
+#####################################
+# Packages & Dependencies
+#####################################
+import torch
+from torch import nn
+
+
+#####################################
+# VGG Model Class
+#####################################
+class VGGBlock(nn.Module):
+    '''
+    Defines a modified block in the VGG architecture, 
+    which includes batch normalization between convolutional layers and ReLU activations.
+    
+    Reference: https://poloclub.github.io/cnn-explainer/
+    Reference: https://d2l.ai/chapter_convolutional-modern/vgg.html
+    
+    Args:
+        num_convs (int): Number of consecutive convolutional layers + ReLU activations.
+        in_channels (int): Number of channels in the input.
+        hidden_channels (int): Number of hidden channels between convolutional layers. 
+        out_channels (int): Number of channels in the output.
+    '''
+    def __init__(self, 
+                 num_convs: int, 
+                 in_channels: int, 
+                 hidden_channels: int,
+                 out_channels: int):
+        super().__init__()
+        
+        self.layers = []
+        
+        for i in range(num_convs):
+            conv_in = in_channels if i == 0 else hidden_channels
+            conv_out = out_channels if i == num_convs-1 else hidden_channels
+            
+            self.layers += [
+                nn.Conv2d(conv_in, conv_out, kernel_size = 3, stride = 1, padding = 1),
+                nn.BatchNorm2d(conv_out),
+                nn.ReLU()
+            ]
+    
+        self.layers.append(nn.MaxPool2d(kernel_size = 2, stride = 2))
+        
+        self.vgg_blk = nn.Sequential(*self.layers)
+            
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        '''
+        Forward pass of VGG block.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, out_channels, new_height, new_width)
+        '''
+        
+        return self.vgg_blk(X)
+    
+class TinyVGG(nn.Module):
+    '''
+    Creates a simplified version of a VGG model, adapted from 
+    https://github.com/poloclub/cnn-explainer/blob/master/tiny-vgg/tiny-vgg.py.
+    The main difference is that the hidden dimensions and number of convolutional layers 
+    remain the same across VGG blocks and the classifier's linear layers has output fewer features.
+    
+    Args:
+        num_blks (int): Number of VGG blocks to put in the model
+        num_convs (int): Number of consecutive convolutional layers + ReLU activations in each VGG block.
+        in_channels (int): Number of channels in the input.
+        hidden_channels (int): Number of hidden channels between convolutional layers. 
+        num_classes (int): Number of class labels.
+        
+    '''
+    def __init__(self, 
+                 num_blks: int, 
+                 num_convs: int, 
+                 in_channels: int, 
+                 hidden_channels: int, 
+                 num_classes: int):
+        super().__init__()
+        
+        self.all_blks = []
+        for i in range(num_blks):
+            conv_in = in_channels if i == 0 else hidden_channels
+            self.all_blks.append(
+                VGGBlock(num_convs, conv_in, hidden_channels, hidden_channels)
+            )
+        
+        self.vgg_body = nn.Sequential(*self.all_blks)
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.LazyLinear(4096), nn.ReLU(), nn.Dropout(0.5),
+            nn.LazyLinear(2048), nn.ReLU(), nn.Dropout(0.5),
+            nn.LazyLinear(num_classes)
+        )
+        
+        self.vgg_body.apply(self._custom_init)
+        self.classifier.apply(self._custom_init)
+        
+    def _custom_init(self, module):
+        '''
+        Initializes convolutional layer weights with Xavier initialization method.
+        Initializes convolutional layer biases to zero.
+        '''
+        if isinstance(module, (nn.Conv2d)):
+            nn.init.xavier_uniform_(module.weight)
+            nn.init.zeros_(module.bias)
+            
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        '''
+        Forward pass of the TinyVGG model.
+        
+        Args:
+            X (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).
+            
+        Returns:
+            torch.Tensor: Logits of shape (batch_size, num_classes).
+        '''
+        
+        X = self.vgg_body(X)
+        return self.classifier(X)

model_training/run_training.py

@@ -0,0 +1,90 @@
+#####################################
+# Packages & Dependencies
+#####################################
+import argparse
+import torch
+from torch import nn
+
+import utils, data_setup, model, engine
+import yaml
+
+# Setup random seeds
+utils.set_seed(0)
+
+# Setup hyperparameters
+parser = argparse.ArgumentParser(fromfile_prefix_chars = '@')
+
+parser.add_argument('-nw', '--num-workers', help = 'Number of workers for dataloaders.',
+                    type = int, default = 0)
+parser.add_argument('-ne', '--num-epochs', help = 'Number of epochs to train model for.', 
+                    type = int, default = 15)
+parser.add_argument('-bs', '--batch-size', help = 'Size of batches to split training set.',
+                    type = int, default = 100)
+parser.add_argument('-lr', '--learning-rate', help = 'Learning rate for the optimizer.', 
+                    type = float, default = 0.001)
+parser.add_argument('-p', '--patience', help = 'Number of epochs to wait before early stopping.', 
+                    type = int, default = 5)
+parser.add_argument('-md', '--min-delta', help = 'Minimum decrease in loss to reset patience.', 
+                    type = float, default = 0.001)
+
+args = parser.parse_args()
+
+
+#####################################
+# Training Code
+#####################################
+if __name__ == '__main__':
+
+    print(f'{'#' * 50}\n'
+          f'\033[1mTraining hyperparameters:\033[0m \n'
+          f'    - num-workers:   {args.num_workers} \n'
+          f'    - num-epochs:    {args.num_epochs} \n'
+          f'    - batch-size:    {args.batch_size} \n'
+          f'    - learning-rate: {args.learning_rate} \n'
+          f'    - patience:      {args.patience} \n'
+          f'    - min-delta:     {args.min_delta} \n'
+          f'{'#' * 50}')
+
+    # Get dataloaders
+    train_dl, test_dl = data_setup.get_dataloaders(root = './mnist_data',
+                                                   batch_size = args.batch_size,
+                                                   num_workers = args.num_workers)
+    
+    # Set up saving directory and file name
+    save_dir = '../saved_models'
+
+    base_name = 'tiny_vgg'
+    mod_name = f'{base_name}_model.pth'
+
+    # Get TinyVGG model
+    mod_kwargs = {
+        'num_blks': 2,
+        'num_convs': 2,
+        'in_channels': 1,
+        'hidden_channels': 10,
+        'num_classes': len(train_dl.dataset.classes)
+    }
+
+    vgg_mod = model.TinyVGG(**mod_kwargs).to(utils.DEVICE)
+
+    # Save model kwargs and train settings
+    with open(f'{save_dir}/{base_name}_settings.yaml', 'w') as f:
+        yaml.dump({'train_kwargs': vars(args), 'mod_kwargs': mod_kwargs}, f)
+
+    # Get loss function and optimizer
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(params = vgg_mod.parameters(), lr = args.learning_rate)
+
+    # Train model
+    mod_res = engine.train(model = vgg_mod,
+                           train_dl = train_dl,
+                           test_dl = test_dl,
+                           loss_fn = loss_fn,
+                           optimizer = optimizer,
+                           num_epochs = args.num_epochs,
+                           patience = args.patience,
+                           min_delta = args.min_delta,
+                           device = utils.DEVICE,
+                           save_mod = True,
+                           save_dir = save_dir,
+                           mod_name = mod_name)

model_training/utils.py

@@ -0,0 +1,64 @@
+#####################################
+# Packages & Dependencies
+#####################################
+import torch
+import random
+import numpy as np
+import os
+
+# Setup device and multiprocessing context
+if torch.cuda.is_available():
+    DEVICE = torch.device('cuda')
+    MP_CONTEXT = None
+elif torch.backends.mps.is_available():
+    DEVICE = torch.device('mps')
+    MP_CONTEXT = 'forkserver'
+else:
+    DEVICE = torch.device('cpu')
+    MP_CONTEXT = None
+
+
+#####################################
+# Functions
+#####################################
+def set_seed(seed: int = 0):
+    '''
+    Sets random seed and deterministic settings for reproducibility across:
+        - PyTorch
+        - NumPy
+        - Python's random module
+    
+    Args:
+        seed (int): The seed value to set.
+    '''
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    
+    torch.use_deterministic_algorithms(True)
+
+def save_model(model: torch.nn.Module,
+               save_dir: str, 
+               mod_name: str):
+    '''
+    Saves the `state_dict()` of a model to the directory 'save_dir.'
+
+    Args:
+        model (torch.nn.Module): The PyTorch model whose state dict and keyword arguments will be saved.
+        save_dir (str): Directory to save the model to.
+        mod_name (str): Filename for the saved model. If this doesn't end with '.pth' or '.pt,' it will be added on for the state_dict.
+
+    '''
+    # Create directory if it doesn't exist
+    os.makedirs(save_dir, exist_ok = True)
+    
+    # Add .pth if it is not in mod_name
+    if not mod_name.endswith('.pth') and not mod_name.endswith('.pt'):
+        mod_name += '.pth'
+
+    # Create save path
+    save_path = os.path.join(save_dir, mod_name)
+
+    # Save model's state dict
+    torch.save(obj = model.state_dict(), f = save_path)

requirements.txt

@@ -1,6 +1,7 @@
-panel
-jupyter
-transformers
-numpy
-torch
-aiohttp
+numpy==2.2.4
+matplotlib==3.10.1
+panel==1.4.5
+param==2.1.1
+plotly==6.0.1
+torch==2.6.0
+torchvision==0.21.0

saved_models/tiny_vgg_less_compute_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94a16b55d2a65b58c30bcad6dcee77d7e45e15221577795aa4de97a508fddced
+size 38494248

saved_models/tiny_vgg_less_compute_settings.yaml

@@ -0,0 +1,13 @@
+mod_kwargs:
+  hidden_channels: 6
+  in_channels: 1
+  num_blks: 2
+  num_classes: 10
+  num_convs: 2
+train_kwargs:
+  batch_size: 100
+  learning_rate: 0.001
+  min_delta: 0.0005
+  num_epochs: 50
+  num_workers: 0
+  patience: 10

saved_models/tiny_vgg_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10b1913e0c2c44d5a76624196371ae83381a13a939afe9e9e9146354206997e9
+size 41711994

saved_models/tiny_vgg_settings.yaml

@@ -0,0 +1,13 @@
+mod_kwargs:
+  hidden_channels: 10
+  in_channels: 1
+  num_blks: 2
+  num_classes: 10
+  num_convs: 2
+train_kwargs:
+  batch_size: 100
+  learning_rate: 0.001
+  min_delta: 0.001
+  num_epochs: 25
+  num_workers: 0
+  patience: 5