PyTorch | 计算高阶导数

#include <iostream>
#include <fstream>
#include <vector>
#include <torch/torch.h>

class Net : public torch::nn::Module
{
public:
    at::Tensor forward(at::Tensor x, at::Tensor y) 
    { 
        return torch::sin(x) * torch::cos(y); 
    }
};

int main(int argc, char* argv[])
{
    // 构造输入数据
    const int numPointsPerDim = 101;
    std::vector<float> x;
    std::vector<float> y;
    for (int i = 0; i < numPointsPerDim; ++i)
    {
        for (int j = 0; j < numPointsPerDim; ++j)
        {
            x.push_back((i - numPointsPerDim / 2) * 0.1);
            y.push_back((j - numPointsPerDim / 2) * 0.1);
        }
    }
    // 转成张量
    torch::Tensor xT =
        torch::from_blob(x.data(), {numPointsPerDim, numPointsPerDim}, torch::kFloat)
            .requires_grad_(true);
    torch::Tensor yT =
        torch::from_blob(y.data(), {numPointsPerDim, numPointsPerDim}, torch::kFloat)
            .requires_grad_(true);
    torch::Tensor weight = torch::ones_like(xT);

    std::shared_ptr<Net> net = std::make_shared<Net>();

    // 函数值 z = sin(x) * cos(y) 
    torch::Tensor z = net->forward(xT, yT);

    // 一阶导数
    torch::Tensor dzdx = torch::autograd::grad({z}, {xT}, {weight}, true, true, false)[0];
    torch::Tensor dzdy = torch::autograd::grad({z}, {yT}, {weight}, true, true, false)[0];

    // 二阶导数
    torch::Tensor d2zdx2 = torch::autograd::grad({dzdx}, {xT}, {weight}, true, true, false)[0];
    torch::Tensor d2zdy2 = torch::autograd::grad({dzdy}, {yT}, {weight}, true, true, false)[0];
    torch::Tensor d2zdxy = torch::autograd::grad({dzdx}, {yT}, {weight}, true, true, false)[0];
    torch::Tensor d2zdyx = torch::autograd::grad({dzdy}, {xT}, {weight}, true, true, false)[0];

    // 输出数据
    std::ofstream os;
    os.open("results.csv");
    for (int i = 0; i < numPointsPerDim; ++i)
    {
        for (int j = 0; j < numPointsPerDim; ++j)
        {
            const float posiX = xT[i][j].item<float>();
            const float posiY = yT[i][j].item<float>();

            const float sinX = std::sin(posiX);
            const float cosX = std::cos(posiX);
            const float sinY = std::sin(posiY);
            const float cosY = std::cos(posiY);
            //
            os << posiX << ',' << posiY << ',';
            os << z[i][j].item<float>() << ',' << sinX * cosY << ',';

            os << dzdx[i][j].item<float>() << ',' << cosX * cosY << ',';
            os << dzdy[i][j].item<float>() << ',' << sinX * -sinY << ',';

            os << d2zdx2[i][j].item<float>() << ',' << -sinX * cosY << ',';
            os << d2zdy2[i][j].item<float>() << ',' << sinX * -cosY << ',';
            os << d2zdxy[i][j].item<float>() << ',' << cosX * -sinY << ',';
            os << d2zdyx[i][j].item<float>() << ',' << cosX * -sinY;
            //
            os << std::endl;
        }
    }
    os.close();

    return 0;
}

cmake_minimum_required(VERSION 3.5)

project(autograd LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

find_package(Torch REQUIRED PATHS "/opt/libtorch/share/cmake/Torch")
include_directories( ${TORCH_INCLUDE_DIRS} )

add_executable(autograd main.cpp)
target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES})

  
.

        
plot.py
#-*- coding:utf-8 -*-

import matplotlib.pyplot as plt 
import numpy as np

titles = ["z_numerical","z_analytical",
          "dzdx_numerical", "dzdx_analytical",
          "dzdy_numerical", "dzdy_analytical",
          "d2zdx2_numerical", "d2zdx2_analytical",
          "d2zdy2_numerical", "d2zdy2_analytical",
          "d2zdxy_numerical", "d2zdxy_analytical",
          "d2zdyx_numerical", "d2zdyx_analytical",
          ]

fig = plt.figure( figsize=(20,20) ) 
ax = plt.axes(projection='3d')

data = np.loadtxt("./results.csv", delimiter=",")

x = data[::,0]
y = data[::,1]
X = x.reshape(101,101)
Y = y.reshape(101,101) 

for idx in range(2,16):
    z = data[::,idx]
    Z = z.reshape(101,101)
    ax.plot_surface(X,Y,Z,cmap='rainbow') 

    plt.title(titles[idx-2],fontsize=35)

    plt.savefig("FIGURE_"+str(idx+1)+".png", bbox_inches='tight')
    plt.cla()