如何构造W变量以基于PyTorch和Plotly绘制二次损失曲面？

现有二次损失公式$z=(1/2)||Aw-b||^2$，其中$A$为4×2矩阵，$w=[x,y]$为2维向量，$b$为4维向量。需绘制$z$关于$x$、$y$的曲面，计划使用Plotly库实现，并借助PyTorch的torch.norm函数计算范数。参考3D曲面绘制示例后编写了如下代码框架，请问应如何修改包含$x$、$y$的$W$变量以完成曲面绘制？

import plotly.graph_objects as go
import numpy as np

A = torch.tensor([[ 0.1542, -0.0682],
        [ 0.8631,  0.6762],
        [-1.4002,  1.1773],
        [ 0.4614,  0.2431]])

b = torch.tensor([-0.2332, -0.7453,  0.9061,  1.2118])

x = np.arange(-1,1,.01)
y = np.arange(-1,1,.01)
X,Y = np.meshgrid(x,y)

W = ??????

Z = 0.5*torch.norm(torch.matmul(A, W)-b)**2

fig = go.Figure(
    data=[go.Surface(z=Z, x=x, y=y, colorscale="Reds", opacity=0.5)])
fig.update_layout(
    title='My title', 
    autosize=False,
    width=500, 
    height=500,
    margin=dict(l=65, r=50, b=65, t=90), 
    scene_aspectmode='cube'
)
fig.show()

核心修改逻辑 #

要让矩阵乘法torch.matmul(A, W)合法，A是4×2矩阵，W需要构造为2×N的张量（N为网格点总数），这样运算后得到4×N的结果，才能和4维的b做广播运算，进而逐点计算损失值。

完整修改代码 #

import plotly.graph_objects as go
import numpy as np
import torch  # 补上缺失的PyTorch导入

A = torch.tensor([[ 0.1542, -0.0682],
        [ 0.8631,  0.6762],
        [-1.4002,  1.1773],
        [ 0.4614,  0.2431]])

b = torch.tensor([-0.2332, -0.7453,  0.9061,  1.2118])

x = np.arange(-1,1,.01)
y = np.arange(-1,1,.01)
X,Y = np.meshgrid(x,y)

# 构造符合维度要求的W张量
x_tensor = torch.tensor(X.flatten(), dtype=torch.float32)
y_tensor = torch.tensor(Y.flatten(), dtype=torch.float32)
W = torch.stack([x_tensor, y_tensor])

# 逐点计算损失并重塑为网格形状
AW_minus_b = torch.matmul(A, W) - b.unsqueeze(1)  # 给b增加维度实现广播
Z = 0.5 * torch.norm(AW_minus_b, dim=0)**2
Z = Z.reshape(X.shape).numpy()  # 转为numpy数组适配Plotly

fig = go.Figure(
    data=[go.Surface(z=Z, x=x, y=y, colorscale="Reds", opacity=0.5)])
fig.update_layout(
    title='二次损失曲面', 
    autosize=False,
    width=500, 
    height=500,
    margin=dict(l=65, r=50, b=65, t=90), 
    scene_aspectmode='cube'
)
fig.show()

关键细节说明 #

• b.unsqueeze(1)：将b从(4,)的一维张量转为(4,1)的二维张量，确保能和(4,N)的AW_minus_b进行逐元素减法运算。
• torch.norm(..., dim=0)：指定对每一列（对应一个网格点的计算结果）计算范数，得到所有网格点的损失值数组。
• Z.reshape(X.shape)：把一维的损失值数组恢复为和X/Y一致的二维网格形状，Plotly才能正确渲染曲面。
• 最后将Z转为numpy数组，因为Plotly的Surface组件对numpy数据兼容性更好。

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