pytorch - 有什么方法可以在pytorch中为张量生成分段函数吗？

Question

我想在 pytorch 中为张量获得这样的分段函数。但我不知道如何定义它。我使用了一种非常愚蠢的方法来做到这一点，但它似乎在我的代码中不起作用。

    def trapezoid(self, X):
        Y = torch.zeros(X.shape)
        Y[X % (2 * pi) < (0.5 * pi)] = (X[X % (2 * pi) < (0.5 * pi)] % (2 * pi)) * 2 / pi
        Y[(X % (2 * pi) >= (0.5 * pi)) & (X % (2 * pi) < 1.5 * pi)] = 1.0
        Y[X % (2 * pi) >= (1.5 * pi)] = (X[X % (2 * pi) >= (1.5 * pi)] % (2 * pi)) * (-2 / pi) + 4
        return Y

你能帮我看看如何设计函数trapezoid，这样对于张量X，我可以直接用trapezoid(X )?

Answer 1

由于您的函数的周期为 2π，我们可以关注 [0,2π]。由于它是分段线性的，因此可以将其表示为 [0,2π] 上的迷你 ReLU 网络，由下式给出:

梯形(x) = 1 - relu(x-1.5π)/0.5π - relu(0.5π-x)/0.5π

因此，我们可以像这样在 Pytorch 中编写整个函数:

import torch
import torch.nn.functional as F
from torch import tensor 
from math import pi

def trapezoid(X):
  # Left corner position, right corner position, height
  a, b, h = tensor(0.5*pi), tensor(1.5*pi), tensor(1.0)

  # Take remainder mod 2*pi for periodicity
  X = torch.remainder(X,2*pi)

  return h - F.relu(X-b)/a - F.relu(a-X)/a

绘图以仔细检查生成正确的图片:

import matplotlib.pyplot as plt

X = torch.linspace(-10,10,1000)
Y = trapezoid(X)
plt.plot(X,Y)
plt.title('Pytorch Trapezoid Function')