如何在PyTorch中实现带约束的优化？

针对优化过程中需要强制参数落在合法空间的需求，最稳定且适配PyTorch原生生态的方案是无约束重参数化，不需要引入额外的带约束优化器，也不会有数值不稳定问题。

核心实现逻辑 #

你不需要直接把带约束的参数（比如必须为正的sigma）设为可训练参数，转而训练一个取值无任何限制的原始参数，每次前向计算损失时，通过单调可导的映射函数把无约束原始参数转换为符合约束要求的实际参数即可。
针对sigma必须为正的场景，直接用指数函数做映射就可以：

• 设可训练的无约束参数为raw_sigma，可取任意实数值
• 实际参与损失计算的sigma为 sigma = torch.exp(raw_sigma)，无论raw_sigma更新到什么值，sigma永远大于0
• 整个映射全程可导，梯度可以正常回传，不会出现硬截断带来的梯度消失、优化卡死问题

不要用优化器step后手动clamp参数的硬约束方案：如果参数更新后被强行截断到边界，截断位置的梯度会完全丢失，大概率会让优化过程卡在边界无法收敛到最优值。

修改后的可运行代码 #

原有代码只需要调整参数定义、损失计算两部分即可，以下版本清理了未使用的冗余导入，同时适配重参数化后的梯度尺度调整了学习率，收敛速度会快很多：

import numpy as np
import torch

# 生成模拟样本
np.random.seed(3)
sample = np.random.normal(loc=5, scale=2, size=(1000, 1))

# 张量定义
X = torch.tensor(sample, dtype=torch.float64, requires_grad=False)
mu = torch.tensor(np.array([0.5]), dtype=torch.float64, requires_grad=True)  # 均值无范围约束，直接作为可训练参数
raw_sigma = torch.tensor(np.array([np.log(5)]), dtype=torch.float64, requires_grad=True)  # 无约束原始参数，初始值对应sigma=5

# SGD优化器
learning_rate = 0.01
optimizer = torch.optim.SGD([mu, raw_sigma], lr=learning_rate)

# 优化循环
for step in range(2000):
    # 无约束参数转合法的正sigma
    sigma = torch.exp(raw_sigma)
    # 计算负对数似然
    nll = X.shape[0] * torch.log(sigma) + ((((X - mu) / sigma) ** 2) / 2).sum()
    
    optimizer.zero_grad()
    nll.backward()
    
    if step % 100 == 0:
        print(f"负对数似然: {nll.item():.2f} | 估计mu: {mu.item():.3f} | 估计sigma: {sigma.item():.3f}")
    
    optimizer.step()

print(f"\nmu、sigma真实值: 5, 2")

其他约束场景扩展 #

这套重参数化思路可以覆盖绝大多数常见的参数约束需求：

• 参数需要大于固定下界lower_bound：param = lower_bound + torch.exp(raw_param)
• 参数需要落在开区间(a, b)内：param = a + (b - a) * torch.sigmoid(raw_param)
• 参数需要满足概率单纯形约束（和为1、每个值大于0）：用torch.softmax(raw_param, dim=...)做映射即可

内容的提问来源于stack exchange，提问作者henriquebsena