Overview
Concrete tool for defining a DenseNet-161 image classifier model for TorchServe model archiver packaging provided by the KServe sample code.
Description
ImageClassifier extends torchvision.models.densenet.DenseNet with the specific configuration parameters for DenseNet-161: growth_rate=48, block_config=(6, 12, 36, 24), and num_init_features=96. It also overrides load_state_dict() to handle legacy checkpoint key naming conventions where layer names contain dots (e.g., "norm.1", "relu.1", "conv.1") that are no longer allowed in modern PyTorch module names. The override uses a regex pattern to find and rename such keys by removing the extra dots before loading the state dictionary.
Usage
Use this class as the model definition file when packaging a DenseNet-161 model into a MAR archive with torch-model-archiver for deployment on TorchServe via KServe.
Code Reference
Source Location
Signature
class ImageClassifier(DenseNet):
def __init__(self):
...
def load_state_dict(self, state_dict, strict=True):
...
Import
from model import ImageClassifier
I/O Contract
Inputs
Constructor
| Name |
Type |
Required |
Description
|
| (none) |
-- |
-- |
No arguments; DenseNet-161 configuration is hardcoded (growth_rate=48, block_config=(6,12,36,24), num_init_features=96)
|
load_state_dict()
| Name |
Type |
Required |
Description
|
| state_dict |
dict |
Yes |
Model state dictionary, potentially with legacy dot-separated key names
|
| strict |
bool |
No |
Whether to enforce strict matching of state_dict keys (default: True)
|
Outputs
load_state_dict()
| Name |
Type |
Description
|
| result |
NamedTuple |
Result from the parent DenseNet.load_state_dict() with missing_keys and unexpected_keys
|
Usage Examples
Basic Usage
from model import ImageClassifier
import torch
# Create the DenseNet-161 model
model = ImageClassifier()
# Load pretrained weights (handles legacy key naming)
state_dict = torch.load("densenet161-8d451a50.pth")
model.load_state_dict(state_dict)
model.eval()
TorchServe Model Archiver
# Package with torch-model-archiver:
# torch-model-archiver --model-name densenet_161 \
# --version 1.0 \
# --model-file model.py \
# --serialized-file densenet161-8d451a50.pth \
# --handler image_classifier \
# --extra-files index_to_name.json
Related Pages
Page Connections
Double-click a node to navigate. Hold to expand connections.
