Extract and Inspect State Dict

easyFREEModel SerializationPyTorch Advanced

Problem Statement

Extract a model's state dictionary and inspect its contents.

Background

model.state_dict() returns an OrderedDict mapping parameter names to tensors. This is the standard way to save and inspect model weights.

Your Task

The starter code creates a model with named layers. Extract the model's state dictionary to inspect its parameter names and shapes.

The inspection code (listing keys, shapes) is pre-filled.

Output Format

Returns a dictionary with "keys", "num_params", "fc1_weight_shape", and "fc2_bias_shape".

Example

Input

None

Output

{'keys': ['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'], 'num_params': 4, 'fc1_weight_shape': [8, 4], 'fc2_bias_shape': [2]}

Reasoning

Input: the function takes no arguments. It builds a two-layer network: fc1 maps 4 features to 8, then fc2 maps 8 features to 2.
Math: a fully connected layer from $n_{in}$ to $n_{out}$ features computes $y = Wx + b$ , with $W \in \mathbb{R}^{n_{out} \times n_{in}}$ and $b \in \mathbb{R}^{n_{out}}$ . So fc1 holds an $8 \times 4$ weight and an 8-vector bias, and fc2 holds a $2 \times 8$ weight and a 2-vector bias.
Keys: each learnable tensor is named layer name, a dot, then its role. In the order the layers were added, that gives ['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'].
Output: num_params is 4 because it counts tensors, not scalars. The scalar count would be $32 + 8 + 16 + 2 = 58$ . fc1_weight_shape is [8, 4] and fc2_bias_shape is [2].
Notice: the weight shape is (outputs, inputs), not (inputs, outputs), which is why fc1 reads [8, 4] and not $[4, 8]$ . The random seed changes only the values inside the tensors. The names and shapes stay the same.

Constraints

Use model.state_dict()
Inspect keys and shapes
Use named modules

Learn this in

PyTorch FundamentalsFoundations

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primer

Background Knowledge

Introduction to PyTorch Models

PyTorch is a popular deep learning framework that provides a dynamic computation graph and automatic differentiation. In PyTorch, a model is essentially a composition of modules, where each module represents a neural network component, such as a fully connected layer (nn.Linear), a convolutional layer (nn.Conv2d), or a recurrent neural network (nn.RNN). The nn.Sequential container is used to create a sequence of modules, allowing for easy construction of complex neural network architectures.

Model Serialization and State Dictionaries

Model serialization refers to the process of saving and loading a model's parameters, allowing for the preservation of trained models and their deployment in various applications. In PyTorch, model serialization is achieved through the use of state dictionaries, which are ordered dictionaries that map parameter names to tensors. The model.state_dict() method returns the state dictionary of a model, providing a way to inspect and manipulate the model's parameters. State dictionaries are essential for tasks such as model checkpointing, fine-tuning, and knowledge distillation.

Understanding Model Parameters

A model's parameters are the learnable components that are adjusted during training to minimize the loss function. In the context of neural networks, parameters typically include weights and biases for each layer. The shape and size of these parameters depend on the specific architecture and layer configurations. For example, the weights of a fully connected layer (nn.Linear) are represented as a 2D tensor, while the biases are 1D tensors. Understanding the structure and organization of model parameters is crucial for tasks such as model inspection, debugging, and optimization.

Algorithm/Approach

The general approach to solving this problem involves creating a PyTorch model using nn.Sequential, extracting its state dictionary using model.state_dict(), and then inspecting the contents of the state dictionary to extract the required information. This involves understanding the structure of the state dictionary, navigating its keys and values, and applying basic tensor operations to extract the desired information.

Step-by-Step Strategy

To solve this problem, follow these steps:

Import the necessary PyTorch modules and set the random seed using torch.manual_seed(42).
Create an nn.Sequential model with the specified named modules ("fc1" and "fc2").
Extract the state dictionary of the model using model.state_dict().
Inspect the state dictionary to extract the required information, including:
- The list of all keys in the state dictionary.
- The number of parameter tensors.
- The shape of fc1.weight as a list.
- The shape of fc2.bias as a list.
Organize the extracted information into a dictionary with the specified keys ("keys", "num_params", "fc1_weight_shape", and "fc2_bias_shape").

Common Pitfalls

When implementing this solution, watch out for the following common pitfalls:

Forgetting to set the random seed, which can lead to inconsistent results.
Incorrectly navigating the state dictionary, which can result in errors or incorrect information.
Failing to convert tensor shapes to lists, which can cause type mismatches.

Time & Space Complexity

The time complexity of this solution is O(1), as it involves a constant number of operations, including creating the model, extracting the state dictionary, and inspecting its contents. The space complexity is also O(1), as the solution only requires a fixed amount of memory to store the model, state dictionary, and extracted information.

solution.py

Python 3.13

Test results