Initialize Only Specific Layer Types

mediumPROWeight InitializationPyTorch Advanced

Problem Statement

Apply different initialization strategies to different layer types in a model.

Background

Different layer types benefit from different initialization. Linear layers with ReLU need Kaiming init, while BatchNorm layers are typically initialized to scale=1, shift=0.

Your Task

The starter code creates a model with Linear and BatchNorm layers. Define an initialization function that applies appropriate initialization for each layer type: Kaiming for Linear layers and standard init for BatchNorm layers. Apply it to the model.

The verification code is pre-filled.

Output Format

Returns a dictionary with "linear_bias_sum", "bn_weight_all_ones", "bn_bias_all_zeros", "num_linear", and "num_bn".

Example

Input

None

Output

{'linear_bias_sum': 0.0, 'bn_weight_all_ones': True, 'bn_bias_all_zeros': True, 'num_linear': 3, 'num_bn': 2}

Reasoning

Input: none. The model is fixed: three Linear layers ( $4 \to 8$ , $8 \to 4$ , $4 \to 2$ ) with BatchNorm layers over $8$ and $4$ features between them, so num_linear is $3$ and num_bn is $2$ .
Linear layers: Kaiming (He) init for ReLU draws each weight from $\mathcal{N}(0, \sigma^2)$ with $\sigma = \sqrt{2 / \text{fan-in}}$ , i.e. $\sqrt{2/4} \approx 0.707$ for the first and last layers and $\sqrt{2/8} = 0.5$ for the middle one. Biases are $0$ .
Bias sum: there are $8 + 4 + 2 = 14$ bias values, all $0$ , so linear_bias_sum is $0.0$ .
BatchNorm layers: the scale $\gamma$ is $1$ and the shift $\beta$ is $0$ for all $8 + 4 = 12$ features, so bn_weight_all_ones and bn_bias_all_zeros are both True.
Notice: a freshly built BatchNorm already has $\gamma = 1$ , $\beta = 0$ , so those flags are True even without init. The bias sum is the real check: default Linear biases are random and sum to about $0.6007$ here, so $0.0$ proves your init ran.

Constraints

Different init for Linear vs BatchNorm
Linear: kaiming_normal_ weights, zeros_ bias
BatchNorm: ones_ weight, zeros_ bias

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 primerPREVIEW

Background Knowledge

The problem revolves around weight initialization in deep learning models, specifically using PyTorch. Weight initialization is a crucial step in training neural networks, as it can significantly impact the convergence speed and overall performance of the model. Different initialization strategies can be applied to different types of layers, such as linear layers and batch normalization layers. The Kaiming normal initialization, for example, is often used for linear layers, while batch normalization layers typically use a different initialization scheme.

In PyTorch, the nn module provides various layer types, including nn.Linear and nn.BatchNorm1d. Each layer has its own set of weights and biases, which need to be initialized before training the model. The torch.nn.init module offers several initialization functions, such as kaiming_normal_ and constant_, which can be used to initialize the weights and biases of different layers. Understanding how to apply these initialization functions to specific layer types is essential for solving this problem.

The problem also requires working with PyTorch's tensor operations, such as summing and comparing tensor values. The torch.manual_seed function is used to set the random seed for reproducibility. Additionally, the problem involves creating a model with multiple layers and applying different initialization strategies to each layer type. This requires a good understanding of PyTorch's module and tensor operations.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Create a PyTorch model with the specified layers
Apply different initialization strategies to each layer type
Extract and process the initialized weights and biases
Compute the required metrics, such as the sum of linear biases and the values of batch normalization weights and biases

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results