有什么方法可以为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

您能帮助我了解如何设计函数梯形，以便对于张量 X ，我可以直接使用梯形（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')