pytorch相关语法，Python特性

记录一些pytorch的笔记，包括与c++的底层绑定。已经一些python的语法。

---

性能优化

cProfile显示各部分运行时间，优化时间长，调用重复的代码。

爱因斯坦求和 einsum

口诀

外部重复做乘积
内部重复把数取
从有到无要求和
重复默认要丢弃(ij, jk-> ), 默认丢掉i，变成ik

用循环实现einsum：

import torch
import numpy as np

a = torch.randn(2,3,5,7)
b = torch.randn(11,13,3,17,5)
# p = 2, q = 3, r = 5, s = 7
# t = 11, u = 13, v = 17, r = 5
torch_ein_out = torch.einsum('pqrs,tuqvr->pstuv', [a, b]).numpy()
torch_org_out = torch.tensordot(a, b, dims=([1, 2], [2, 4])).numpy()

np_a = a.numpy()
np_b = b.numpy()
# 循环展开实现
np_out = np.empty((2, 7, 11, 13, 17), dtype=np.float32)
# 自由索引外循环
# 这里就是 p,s,t,u和v
for p in range(0, 2):
    for s in range(0, 7):
        for t in range(0, 11):
            for u in range(0, 13):
                for v in range(0, 17):
                    # 求和索引内循环
                    # 这里是 q和r
                    sum_result = 0
                    for q in range(0, 3):
                        for r in range(0, 5):
                            sum_result += np_a[p, q, r, s] * np_b[t, u, q, v, r]
                    np_out[p, s, t, u, v] = sum_result

print("is np_out == torch_ein_out ?", np.allclose(torch_ein_out, np_out, atol=1e-6))
print("is torch_ein_out == torch_org_out ?", np.allclose(torch_ein_out, torch_org_out, atol=1e-6))

# 终端打印输出
# is np_out == torch_ein_out ? True
# is torch_ein_out == torch_org_out ? True

下图解释了张量相乘，类似GCN中的，AX，但是应用爱因斯坦求和(X,A)。

一文学会 Pytorch 中的 einsum
案例. 四维张量乘三维张量
知乎einsum文章
优雅地实现多头自注意力——使用einsum（爱因斯坦求和）进行矩阵运算
einsum

与c++底层的绑定方法

PYBIND11_MODULE 宏绑定函数

该宏用于定义一个 PyBind11 模块，example 是python模块的名称，m 是模块对象，可以用于绑定类和函数，以便在 Python 中使用。

#include <pybind11/pybind11.h>

int add(int i, int j) {
    return i + j;
}

PYBIND11_MODULE(example, m) {
    m.def("add", &add, "A function which adds two numbers");
}

py::class_ 绑定c++类

pybind11::class_<MyClass>(m, "MyClass")
“MyClass”: 是py中的新类名

: 指的绑定c++的类名

- m.def 是用于将c++全局函数绑定到 Python 模块上的方法。
-    .def 是用于将类的方法绑定到 Python 上的接口。

使用 PyTorch 的 C++ 扩展 (TorchScript)