Machine Learning Academy · レッスン

nn.Moduleによるフィードフォワードネットワークの構築

nn.Moduleを継承し、__init__でLinear層とReLU層を積み重ね、forward処理を実装して、ダミー入力で出力形状を検証します。

「nn.Moduleによるフィードフォワードネットワークの構築」はCoddyKit上の無料Machine Learning Academyレッスンです。これはレッスン3/4です。下記で完全なレッスンを無料で読むことができます。その後、ブラウザ内の組み込みコードエディタと24時間対応のAIチューターでハンズオン演習できます。これはMachine Learning Academy学習パスの一部であり、ウェブとCoddyKitアプリ全体で進捗が同期されます。 Machine Learning Academyコースには全4レッスンが含まれています。

このレッスンの一部はまだ翻訳されておらず、英語で表示されています。

What Is nn.Module?

nn.Module is PyTorch's base class for all neural network components. Every layer, activation function, loss function, and complete model in PyTorch subclasses nn.Module. It provides parameter management, device movement, serialisation, and training/eval mode switching out of the box. By subclassing it you get all this functionality for free and only need to define the architecture and the forward pass.

import torch.nn as nn

# Inspect what nn.Module gives you
model = nn.Linear(4, 2)
print(type(model))               # <class 'torch.nn.modules.linear.Linear'>
print(isinstance(model, nn.Module))  # True
print(list(model.parameters()))  # weight and bias tensors

Subclassing nn.Module: init and forward

Creating a custom network requires two methods. In __init__ you call super().__init__() and define all learnable layers as attributes. In forward you describe how data flows through those layers. PyTorch tracks any nn.Module or nn.Parameter assigned to self as a trainable component. Calling the model like a function (model(x)) automatically invokes forward and runs any registered hooks.

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet(4, 16, 3)
print(model)

Linear Layers: nn.Linear

nn.Linear(in_features, out_features) applies the transformation y = x @ W.T + b where W is a weight matrix and b is a bias vector. The layer automatically initialises weights with kaiming uniform distribution and biases with uniform distribution. It is the building block of feedforward networks, also called fully connected or dense layers.

import torch
import torch.nn as nn

layer = nn.Linear(5, 3)  # 5 inputs, 3 outputs

print('Weight shape:', layer.weight.shape)  # (3, 5)
print('Bias shape:',   layer.bias.shape)    # (3,)

# Forward pass with a batch of 8 samples
x = torch.randn(8, 5)
out = layer(x)
print('Output shape:', out.shape)   # (8, 3)

Activation Functions: ReLU, Sigmoid, Tanh

Activation functions introduce non-linearity, allowing the network to learn complex patterns that a stack of linear layers cannot represent. ReLU (max(0, x)) is the default for hidden layers — fast and avoids vanishing gradients. Sigmoid squashes output to [0, 1], used for binary classification outputs. Tanh squashes to [-1, 1], commonly used in RNNs. All are available as modules in nn.

import torch
import torch.nn as nn

x = torch.tensor([-2.0, -0.5, 0.0, 0.5, 2.0])

print('ReLU:   ', nn.ReLU()(x))
# tensor([0.0, 0.0, 0.0, 0.5, 2.0])

print('Sigmoid:', nn.Sigmoid()(x))
# tensor([0.12, 0.38, 0.50, 0.62, 0.88])

print('Tanh:   ', nn.Tanh()(x))
# tensor([-0.96, -0.46,  0.00,  0.46,  0.96])

Stacking Layers with nn.Sequential

nn.Sequential is a convenient container that passes the output of each module as the input to the next. It is ideal for simple feedforward architectures where data flows linearly. For more complex networks with skip connections, multiple inputs, or branching paths, you need the full nn.Module subclass approach. Sequential networks can still be extended by subclassing and using the Sequential as a sub-block.

import torch
import torch.nn as nn

# Build with nn.Sequential
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

x = torch.randn(16, 10)    # batch of 16
out = model(x)
print(out.shape)            # torch.Size([16, 1])

Verifying Output Shapes with Dummy Input

A critical debugging technique when building networks is to run a dummy tensor through the model before training. This verifies that all layer dimensions are compatible and reveals shape mismatches immediately. The dummy tensor has the same shape as your real data but contains random values — it just exercises the forward path. Always do this after defining a new architecture.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.net(x)

model = MLP()
dummy = torch.randn(32, 784)   # batch of 32 MNIST images
out = model(dummy)
print(out.shape)                # torch.Size([32, 10]) -- correct!

Listing and Counting Parameters

Understanding the total number of trainable parameters in a model gives you a sense of its capacity and memory requirements. model.parameters() returns an iterator over all learnable tensors; model.named_parameters() pairs each tensor with its name for inspection. A compact parameter counter is one of the first utilities every PyTorch practitioner builds.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters()
                if p.requires_grad)

print(f'Total params:     {total_params:,}')   # 203,530
print(f'Trainable params: {trainable:,}')       # 203,530

for name, p in model.named_parameters():
    print(name, p.shape)

Moving the Model to GPU

Moving the model to the GPU is as simple as calling model.to(device). This transfers all parameters and buffers registered in the module to the target device. After this call, all forward pass computations happen on the GPU automatically — as long as your input tensors are also on the same device. Mixing CPU and GPU tensors raises a runtime error.

import torch
import torch.nn as nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2)
).to(device)   # move entire model to device

# Input must be on the same device
x = torch.randn(5, 4).to(device)
out = model(x)
print(out.device)   # cuda:0 (or cpu)

Training vs Eval Mode

Some layers (Dropout, BatchNorm) behave differently during training and inference. Calling model.train() enables stochastic behaviour (random dropout, batch statistics); model.eval() switches to deterministic inference behaviour. Forgetting to call model.eval() before evaluation leads to inconsistent results because Dropout randomly zeros activations. Always pair these calls with torch.no_grad() during inference for maximum efficiency.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.Dropout(p=0.5),
    nn.Linear(8, 2)
)

# Training mode: dropout is active
model.train()
x = torch.randn(4, 4)
print(model(x))   # some activations zeroed randomly

# Eval mode: dropout is disabled
model.eval()
with torch.no_grad():
    print(model(x))   # deterministic output

Custom Network: Multi-Layer Perceptron

Putting it all together: a multi-layer perceptron (MLP) is a feedforward network with one or more hidden layers between input and output. Each hidden layer applies a linear transformation followed by a non-linear activation. The output layer's activation depends on the task — no activation for regression, softmax for multi-class classification, sigmoid for binary classification. The example below builds a 3-layer MLP for 10-class classification.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dims, out_dim):
        super().__init__()
        layers = []
        prev = in_dim
        for h in hidden_dims:
            layers.append(nn.Linear(prev, h))
            layers.append(nn.ReLU())
            prev = h
        layers.append(nn.Linear(prev, out_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)

model = MLP(784, [256, 128, 64], 10)
dummy = torch.randn(32, 784)
print(model(dummy).shape)  # torch.Size([32, 10])

Saving and Loading Network State

Neural network training is expensive, so you save the model after training. PyTorch's convention is to save only the state_dict — a dictionary of parameter tensors — rather than the entire model object. This avoids pickle dependency on the class definition. To restore, create a fresh model instance with the same architecture, then load the state dict with load_state_dict. Always set model.eval() after loading for inference.

import torch
import torch.nn as nn

model = nn.Linear(4, 2)

# Save only parameters (recommended)
torch.save(model.state_dict(), '/tmp/model.pt')

# Restore
new_model = nn.Linear(4, 2)  # same architecture
new_model.load_state_dict(torch.load('/tmp/model.pt'))
new_model.eval()

x = torch.randn(3, 4)
print(new_model(x))   # same output as original model

Quick Check

Test your understanding of Machine Learning with Python concepts from this lesson.

Lesson Recap

In this lesson you learned: nn.Module is the base class for all PyTorch neural networks, requiring __init__ to define layers and forward to define data flow, nn.Sequential provides a simple container for linear stacks of layers, and model.train() / model.eval() switches behaviour of layers like Dropout and BatchNorm. Next up we write the complete training loop including loss, optimizer, and multiple epochs.

よくある質問