PyTorch torch.nn.MSELoss Function

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

torch.nn.MSELossIt is the mean squared error loss function in PyTorch.

It computes the mean of the squared differences between predicted values and target values, commonly used in regression tasks.

Function Definition

torch.nn.MSELoss(reduction='mean')

Parameter Description:

  • reduction(str): Loss aggregation method. Optional'mean'、'sum'、'none'. Default is'mean'。

Mathematical Principle

MSE loss formula:

MSE = (1/n) * Σ(y_pred - y_true)²

Usage Examples

Example 1: Basic Usage

Calculate regression loss:

Example

import torch
import torch.nn as nn

# Create MSE loss
criterion = nn.MSELoss()

# Predicted values and true values
predictions = torch.tensor([3.0, 4.0, 5.0])
targets = torch.tensor([2.8, 4.2, 4.9])

# Calculate loss
loss = criterion(predictions, targets)

print("Predicted values:", predictions.tolist())
print("Target values:", targets.tolist())
print("MSE loss:", loss.item())

# Manual verification
manual_mse = ((predictions - targets) ** 2).mean()
print("Manual calculation:", manual_mse.item())

Example 2: Complete Regression Training

Regression task training process:

Example

import torch
import torch.nn as nn
import torch.optim as optim

# Simple regression model
class RegressionNet(nn.Module):
    def __init__(self):
        super(RegressionNet, self).__init__()
        self.fc = nn.Linear(10, 1)

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

model = RegressionNet()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training data
X = torch.randn(100, 10)
y = torch.randn(100, 1)

# Training steps
model.train()
optimizer.zero_grad()
predictions = model(X)
loss = criterion(predictions, y)
loss.backward()
optimizer.step()

print("Training loss:", loss.item())

Example 3: reduction Parameter

Different aggregation methods:

Example

import torch
import torch.nn as nn

pred = torch.tensor([3.0, 4.0, 5.0, 6.0])
target = torch.tensor([3.5, 3.8, 5.2, 5.5])

# mean: average loss
loss_mean = nn.MSELoss(reduction='mean')(pred, target)

# sum: total loss
loss_sum = nn.MSELoss(reduction='sum')(pred, target)

# none: no aggregation
loss_none = nn.MSELoss(reduction='none')(pred, target)

print("Mean:", loss_mean.item())
print("Sum:", loss_sum.item())
print("None:", loss_none.tolist())

Example 4: Handling Multi-dimensional Output

Multi-dimensional regression:

Example

import torch
import torch.nn as nn

# Multi-dimensional output
pred = torch.randn(4, 3)  # batch=4, 3 outputs
target = torch.randn(4, 3)

criterion = nn.MSELoss()
loss = criterion(pred, target)

print("Predicted shape:", pred.shape)
print("Target shape:", target.shape)
print("MSE loss:", loss.item())

FAQ

Q1: Which is better, MSE or MAE?

  • MSE: Sensitive to outliers, stable gradients
  • MAE (L1): Robust to outliers, unstable gradients

Q2: Can MSE be used if outputs have negative values?

Yes, MSE does not restrict the output range.

Q3: What loss is used for classification tasks?

Use CrossEntropyLoss for classification tasks.

Use Cases

nn.MSELossMain application scenarios include:

  • Regression tasks: House price prediction, numerical estimation
  • Continuous value prediction: Target tracking
  • Generative models: Certain losses for VAE, GAN

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

Other Extensions