6

Module Onboarding Lab

Manuel edited this page 2024-07-03 19:31:14 +02:00

Table of Contents

• Goal
• Components
• Relevant Files
• Next

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

The goal of this lab is to give first-hand experience with defining and testing a new module within PyTorch.

The deliverable will be a PR (that won’t be merged into master) containing all the code for the different sections below. You should add both your mentor as well as @jbschlosser as a reviewer for this PR.

Goal

For this lab, we are going to create a module named nn.Bias that simply adds a learnable vector parameter to the last dimension of its input (i.e. a bias-only analogue of PyTorch's Linear module). The module should accept a constructor argument named num_features that defines the size of its vector parameter, and the vector should be randomly initialized from a standard normal distribution. Example usage for this module should look like:

import torch
from torch import nn

m = nn.Bias(num_features=5)
input = torch.randn(10, 5)
output = m(input)  # output should be equal to input + module's bias parameter

Components

While this is not the case for all modules, it is common for the "meat" of a module to be defined as a single PyTorch operator. On top of this, there are generally 4 wrappers defined that utilize the underlying operator and provide a nice UX for the user, whether they are using the Python or C++ APIs:

• Python API functional form - calls the underlying operator through Python bindings
• Python API module class - calls the Python API functional form
• C++ API functional form - calls the underlying operator
• C++ API module class - calls the C++ API functional form

You should define a PyTorch operator for the "meat" of the module and implement each of the wrappers above.

Relevant Files

• aten/src/ATen/native/native_functions.yaml - add the definition for the underlying PyTorch operator here (see Tensor-and-Operator-Basics)
• aten/src/ATen/native/Linear.cpp - add core logic for the operator here
• torch/nn/modules/linear.py - add the Python module class definition here
• torch/nn/modules/__init__.py - add the module class name to the list in __all__
• torch/nn/functional.py - add the Python functional form of the module here
• torch/csrc/api/include/torch/nn/functional/linear.h - add the C++ functional form here
• torch/csrc/api/include/torch/nn/modules/linear.h / torch/csrc/api/src/nn/modules/linear.cpp - add the C++ module class definitions here
• torch/csrc/api/include/torch/nn/options/linear.h / torch/csrc/api/src/nn/options/linear.cpp - add options for the module here (i.e. num_features)
• torch/testing/_internal/common_nn.py - add a dict entry to the module_tests list defining how to generate generic tests for the new module (note: this part is in the process of being improved)
• test/test_nn.py - add any module functionality tests here as functions on TestNN or TestNNDeviceType (for device-specific tests)
• test/cpp/api/functional.cpp - add C++ API functional tests here
• test/cpp/api/modules.cpp - add C++ API module tests here
• test/cpp_api_parity/parity-tracker.md - add module and functional entries to the parity tracker to check C++ vs. Python behavior (used by build/bin/test_api, which runs the C++ module and functional tests)
• torch/overrides.py - add override lambda here (used by test/test_overrides.py)
• torch/nn/functional.pyi.in - add skeleton signature here (used by mypy)
• ??? - other updates may be needed to pass all CI checks & tests

Next

Unit 6: CUDA basics - CUDA-basics

I would love to contribute to PyTorch!