Plotting a quadratic loss using norm function from Pytorch and showing it using Plotly

user_11748994 0
Viewed 159

I have a quadratic loss z=(1/2)||Aw-b||^2 where A is 4x2 matrix, w=[x,y] is a 2d vector, and b is a 4d vector. If we plot z, there would be a surface in terms of x,y. I want to plot z using Plotly library. To do this, I want to use Pytorch and the function torch.norm for calculating the norm. Here is a worked example for plotting a 3d surface and I want to modify it as follows:

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()

Question:

How should I modify W which includes x,y to plot the surface?

pytorch plotly python
2 Answers

You could simply do:

Z = [[0.5 * torch.norm(torch.matmul(A, torch.tensor([float(xx), float(yy)]))-b)**2 for xx in x] for yy in y]

Update: You can improve the performance significantly by using torch's micro batch feature. For this you have to reshape your data to lists of matrices. That means you have to extend tensor A to a list that contains only one matrix and W to a list that contains all mesh points, each as matrix.

import plotly.graph_objects as go
import torch

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 = torch.arange(-1, 1, 0.01)
y = torch.arange(-1, 1, 0.01)
W = torch.reshape(torch.cartesian_prod(x, y), (len(x) * len(y), 2, 1))

V = torch.reshape(torch.matmul(A, W), (len(x), len(y), 4)) - b
Z = 0.5 * torch.norm(V, dim=2)**2
user_186672250

After a lot of trial and error, turns out the most optimized code is to write the linear algebra operation "by hand" (without using any linear algebra packages at all) and JIT compiling it (with Numba).

import plotly.graph_objects as go
from numba import njit, float32


@njit(float32(float32, float32))
def op(x, y):
    return 0.5 * ((0.1542*x - 0.0682*y - 0.2332)*(0.1542*x - 0.0682*y - 0.2332) +
                  (0.8631*x + 0.6762*y - 0.7453)*(0.8631*x + 0.6762*y - 0.7453) +
                  (-1.4002*x + 1.1773*y + 0.9061)*(-1.4002*x + 1.1773*y + 0.9061) +
                  (0.4614*x + 0.2431*y + 1.2118)*(0.4614*x + 0.2431*y + 1.2118))


interval = [i/100 for i in range(-100, 100)]
z = [[op(xi, yi) for yi in interval] for xi in interval]
fig = go.Figure(data=[go.Surface(z=z, x=interval, y=interval)])
fig.show()
user_117090210
Related