PyTorch中Linear层实现的SOM未训练问题排查

你的SOM权重没有更新的核心原因是拓扑损失（topo_loss）的计算完全没有关联到SOM的权重参数，导致SOM权重的梯度始终为0，反向传播时没有梯度流驱动其更新。

看你的训练代码：

• squared_l2_norm_dists是SOM节点位置（locations）之间的距离平方，和SOM权重model.som.som_wts.weight完全无关
• topo_loss最终由topo_neighb * squared_l2_norm_dists计算而来，整个过程没有用到SOM权重，自然不会对其产生梯度

正确的DESOM拓扑损失需要把latent code z和SOM权重的关联引入，让模型通过损失驱动SOM权重向对应latent code靠拢，同时保留拓扑约束。以下是修正后的核心训练代码：

# Get latent code and reconstruction-
z, x_recon = model(x)

optimizer.zero_grad()

# Autoencoder reconstruction loss-
recon_loss = F.mse_loss(input = x_recon, target = x)

# SOM 训练代码-
# 1. 计算z到所有SOM权重的距离
l2_dist_z_soms = torch.cdist(x1 = z, x2 = model.som.som_wts.weight, p = p_norm)
mindist, bmu_indices = torch.min(l2_dist_z_soms, -1)
bmu_locations = locations[bmu_indices]

# 2. 计算每个SOM节点到BMU的拓扑距离（位置距离）
squared_topo_dists = torch.square(torch.cdist(x1 = locations, x2 = bmu_locations, p = p_norm))

# 3. 计算高斯拓扑邻居权重
global step
curr_sigma_val = sigma_0 * torch.exp(-step / lmbda_val)
step += 1
topo_neighb = torch.exp(-squared_topo_dists / ((2 * torch.square(curr_sigma_val)) + 1e-6))
# 转置以匹配batch维度：[batch_size, num_som_nodes]
topo_neighb = topo_neighb.T

# 4. 计算拓扑损失：让z靠近BMU及其邻居的SOM权重
# 损失 = 拓扑权重 * (z与SOM权重的距离平方)，求和后取batch平均
topo_loss = (topo_neighb * torch.square(l2_dist_z_soms)).sum(1).mean()

# Compute total loss-
total_loss = recon_loss + (gamma * topo_loss)
# gamma = 0.001

# Compute gradients wrt computed loss-
total_loss.backward()
        
# Perform one step of gradient descent-
optimizer.step()

• 确认你的optimizer已经包含了SOM层的参数，比如初始化时用optimizer = torch.optim.Adam(model.parameters(), lr=...)，确保model.som.som_wts.weight在参数列表中
• 训练过程中可以打印model.som.som_wts.weight.grad，如果修复后梯度不为0，说明问题已解决

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