Per-Layer Learning Rates

easyFREECustom OptimizersPyTorch Advanced

Problem Statement

Configure different learning rates for different layers using parameter groups.

Background

PyTorch optimizers accept a list of parameter groups, each with its own hyperparameters. This is essential for fine-tuning where you want lower LR for pretrained layers.

Your Task

The starter code creates a 2-layer model. Create an Adam optimizer with different learning rates per layer: 0.001 for fc1 and 0.01 for fc2.

Output Format

Returns a dictionary with "num_param_groups", "group0_lr", "group1_lr", "group0_num_params", and "group1_num_params".

Example

Input

None

Output

{'num_param_groups': 2, 'group0_lr': 0.001, 'group1_lr': 0.01, 'group0_num_params': 2, 'group1_num_params': 2}

Reasoning

The function takes no arguments. It builds a fixed two-layer network: fc1 maps $4 \to 8$ , a ReLU follows, and fc2 maps $8 \to 2$ . The ReLU has no learnable weights.
The model's learnable tensors are split into two groups, each with its own step size $\eta$ . Group 0 holds fc1 with $\eta_0 = 0.001$ and group 1 holds fc2 with $\eta_1 = 0.01$ . So num_param_groups is 2.
With Adam, every update in a group is scaled by that group's $\eta$ . So the output layer moves $\eta_1 / \eta_0 = 10$ times faster than the first layer.
Each linear layer has two tensors, a weight and a bias. For fc1 they are an $8 \times 4$ matrix (32 numbers) and a length-8 vector. For fc2 they are a $2 \times 8$ matrix (16 numbers) and a length-2 vector.
Notice that group0_num_params and group1_num_params are both 2. They count tensors, not scalars: 40 values and 18 values both come out as 2. The group order also matters, because the first group listed becomes group 0.

Constraints

Two parameter groups with different LRs
fc1: lr=0.001, fc2: lr=0.01
Use Adam optimizer

Learn this in

PyTorch FundamentalsFoundations Fine-tuningLLM Study Plan Text ClassificationNLP Study Plan 3D Gaussian Splatting3D Computer Vision

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Implement Bilinear Form with Einsumhard

Concept primer

Background Knowledge

Introduction to PyTorch Optimizers

PyTorch provides a variety of optimizers that can be used to update the parameters of a model during training. These optimizers include SGD, Adam, RMSprop, among others. Each optimizer has its own set of hyperparameters that can be tuned for better performance. One of the key features of PyTorch optimizers is the ability to accept a list of parameter groups, each with its own set of hyperparameters.

Parameter Groups

Parameter groups are a way to group the parameters of a model into different categories, each with its own set of hyperparameters. This is useful when fine-tuning a pre-trained model, where the pre-trained layers may require a lower learning rate than the newly added layers. By using parameter groups, you can specify different learning rates for different layers of the model. For example, you can specify a lower learning rate for the pre-trained layers and a higher learning rate for the newly added layers.

PyTorch Sequential Model

A Sequential model in PyTorch is a container that holds a sequence of modules. Each module can be a linear layer, a convolutional layer, a recurrent layer, or any other type of layer. The Sequential model is a convenient way to build a model by adding layers one after the other. In this problem, we are using a Sequential model with two linear layers (fc1 and fc2) and a ReLU activation function.

Algorithm/Approach

The general approach to solve this problem is to:

Create a PyTorch model using the Sequential container
Define the parameter groups for the optimizer, where each group has its own set of hyperparameters
Create an optimizer instance with the defined parameter groups
Extract the required information from the optimizer and return it as a dictionary

Step-by-Step Strategy

To implement the solution, follow these steps:

Import the necessary PyTorch modules and set the random seed using torch.manual_seed(42).
Create a Sequential model with the specified layers (fc1, relu, and fc2).
Define the parameter groups for the optimizer. In this case, we have two groups: one for the fc1 parameters and one for the fc2 parameters.
Create an Adam optimizer instance with the defined parameter groups.
Extract the required information from the optimizer, including the number of parameter groups, the learning rate of each group, and the number of parameter tensors in each group.
Return the extracted information as a dictionary.

Common Pitfalls

When implementing the solution, watch out for the following:

Make sure to set the random seed correctly to ensure reproducibility.
Verify that the parameter groups are defined correctly, with the correct learning rates and parameters.
Ensure that the optimizer is created with the correct parameter groups.

Time & Space Complexity

The time complexity of this solution is O(1), since we are only creating a model, defining parameter groups, and extracting information from the optimizer. The space complexity is also O(1), since we are not storing any additional data structures that scale with the input size. The solution only depends on the number of parameter groups and the number of parameters in each group, which is fixed in this problem.

solution.py

Python 3.13

Test results