PyTorch底层原理

┌─────────────────────────────────────────────────────────────┐
│                    Python Layer                             │
│              torch.nn, torch.optim                          │
├─────────────────────────────────────────────────────────────┤
│                    PyTorch C++ Core                         │
│  ┌─────────────────────────────────────────────────────┐   │
│  │                    ATen (Tensor Library)            │   │
│  │              Core tensor operations                 │   │
│  ├─────────────────────────────────────────────────────┤   │
│  │                    C10 (Core Libraries)             │   │
│  │         Device, DataType, Dispatch, Storage         │   │
│  └─────────────────────────────────────────────────────┘   │
├─────────────────────────────────────────────────────────────┤
│                    Backends                                 │
│         CUDA    │    CPU    │    MKLDNN    │    XLA        │
└─────────────────────────────────────────────────────────────┘

# 内存使用监控
import torch
print(torch.cuda.memory_summary())

# 手动清理缓存
torch.cuda.empty_cache()

# 设置内存池大小限制
torch.cuda.set_per_process_memory_fraction(0.8)

# 内存追踪
torch.cuda.memory追踪器()

// my_extension.cpp
#include <torch/extension.h>
#include <cuda.h>

// CUDA Kernel
__global__ void relu_kernel(float* x, float* y, int n) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    if (i < n) y[i] = fmaxf(x[i], 0.0f);
}

// C++ wrapper
torch::Tensor relu_cuda(torch::Tensor x) {
    auto y = torch::empty_like(x);
    int n = x.numel();
    int threads = 256;
    int blocks = (n + threads - 1) / threads;
    relu_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(), 
        y.data_ptr<float>(), 
        n
    );
    return y;
}

// 绑定到Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("relu", &relu_cuda, "ReLU CUDA");
}

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='my_extension',
    ext_modules=[
        CUDAExtension('my_extension', [
            'my_extension.cpp',
        ])
    ],
    cmdclass={'build_ext': BuildExtension}
)

# 自定义Autograd Function
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)
    
    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[x < 0] = 0
        return grad_input