CUDA 在机器学习中的应用 - 直观而全面的解释

在本文中,我们将使用 CUDA 在 GPU 上训练一个 AI 模型,基本上是从头实现 AI,假设几乎没有任何先验知识。

1. 现代计算机的组成

现代计算机主要由以下几个部分组成:

主板:连接所有计算机组件的电路板
CPU(中央处理器):执行计算的核心
RAM(随机存取内存):CPU 的工作内存
GPU(图形处理器):用于加速某些类型的计算

CPU 和 RAM 通常被称为"主机",而 GPU 被称为"设备"。GPU 内部也有自己的处理单元和内存(VRAM)。

2. GPU 在 AI 中的重要性

GPU 相比 CPU 的优势在于:

GPU 有更多但能力较弱的核心,可以进行大规模并行计算
GPU 支持 SIMD(单指令多数据),一个控制电路可以控制多个核心

AI 模型通常涉及大量简单且相对独立的计算,非常适合在 GPU 上并行执行,可以获得 100 倍甚至 1000 倍的速度提升。

3. CUDA 简介

CUDA(Compute Unified Device Architecture)是 NVIDIA 的并行计算平台,允许在 CPU 上运行的应用程序调用 GPU 进行并行数学计算。

CUDA 的核心概念包括:

Kernel:在 GPU 上并行执行的函数
Thread:并行任务的基本单位
Thread Block:共享内存的线程组
Grid:多个 Thread Block 的集合

使用 CUDA 编程的基本流程:

在 CPU 上分配内存并初始化数据
将数据从 CPU 内存复制到 GPU 内存
在 GPU 上执行 kernel 函数
将结果从 GPU 内存复制回 CPU 内存

4. 从头实现神经网络

我们将使用 CUDA 实现一个简单的神经网络,包括以下组件:

Shape 类:用于表示矩阵的形状
NNException 类:用于处理 CUDA 错误
Matrix 类:抽象化 CPU 和 GPU 之间的内存操作
Layer 类:神经网络层的基类
Linear 类:实现全连接层
Sigmoid 和 ReLU 类:实现激活函数
BinaryCrossEntropy 类:实现损失函数

4.1 Shape 类

// shape.hh
#pragma once

struct Shape {
  size_t x, y;
  Shape(size_t x = 1, size_t y = 1);
};

// shape.cu  
#include "shape.hh"

Shape::Shape(size_t x, size_t y) : x(x), y(y) {}

4.2 NNException 类

// nn_exception.hh
#pragma once

#include <exception>
#include <iostream>

class NNException : std::exception {
private:
  const char* exception_message;

public:
  NNException(const char* exception_message);
  virtual const char* what() const throw();
  static void throwIfDeviceErrorsOccurred(const char* exception_message);
};

4.3 Matrix 类

// matrix.hh
#pragma once

#include "shape.hh"
#include "nn_exception.hh"

class Matrix {
private:
  float* device_data; 
  float* host_data;
  Shape shape;

public:
  Matrix(size_t x_dim, size_t y_dim);
  ~Matrix();
  
  void allocateMemory();
  void allocateDeviceMemory();
  void allocateHostMemory();
  
  void copyHostToDevice();
  void copyDeviceToHost();
  
  float& operator[](const int index);
  const float& operator[](const int index) const;

  Shape getShape() const { return shape; }
  float* getDeviceData() { return device_data; }
  float* getHostData() { return host_data; }
};

4.4 Layer 类

// layer.hh
#pragma once

#include "matrix.hh"

class Layer {
public:
  virtual ~Layer() {}
  virtual Matrix& forward(Matrix& input) = 0;
  virtual Matrix& backward(Matrix& output_gradient, float learning_rate) = 0;
};

4.5 Linear 类

// linear.hh
#pragma once

#include "layer.hh"

class Linear : public Layer {
private:
  Matrix weights;
  Matrix bias;
  Matrix input;
  Matrix output;

public:
  Linear(size_t input_size, size_t output_size);
  ~Linear();

  Matrix& forward(Matrix& input) override;
  Matrix& backward(Matrix& output_gradient, float learning_rate) override;
};

4.6 激活函数类

// activation.hh
#pragma once

#include "layer.hh"

class Sigmoid : public Layer {
public:
  Matrix& forward(Matrix& input) override;
  Matrix& backward(Matrix& output_gradient, float learning_rate) override;
};

class ReLU : public Layer {
public:
  Matrix& forward(Matrix& input) override;
  Matrix& backward(Matrix& output_gradient, float learning_rate) override;
};

4.7 BinaryCrossEntropy 类

// binary_cross_entropy.hh
#pragma once

#include "matrix.hh"

class BinaryCrossEntropy {
public:
  static float loss(const Matrix& predictions, const Matrix& targets);
  static Matrix gradient(const Matrix& predictions, const Matrix& targets);
};

5. 训练神经网络

有了以上组件,我们就可以构建一个简单的神经网络并进行训练:

int main() {
  // 准备数据
  Matrix X_train = loadTrainingData();
  Matrix y_train = loadTrainingLabels();
  
  // 构建模型
  Linear layer1(INPUT_SIZE, HIDDEN_SIZE);
  ReLU activation1;
  Linear layer2(HIDDEN_SIZE, OUTPUT_SIZE);
  Sigmoid activation2;
  
  // 训练循环
  for (int epoch = 0; epoch < EPOCHS; epoch++) {
    // 前向传播
    Matrix h1 = layer1.forward(X_train);
    Matrix a1 = activation1.forward(h1);
    Matrix h2 = layer2.forward(a1);
    Matrix y_pred = activation2.forward(h2);
    
    // 计算损失
    float loss = BinaryCrossEntropy::loss(y_pred, y_train);
    
    // 反向传播
    Matrix grad = BinaryCrossEntropy::gradient(y_pred, y_train);
    grad = activation2.backward(grad, LEARNING_RATE);
    grad = layer2.backward(grad, LEARNING_RATE);
    grad = activation1.backward(grad, LEARNING_RATE);
    grad = layer1.backward(grad, LEARNING_RATE);
    
    std::cout << "Epoch " << epoch << ", Loss: " << loss << std::endl;
  }
  
  // 测试模型
  Matrix X_test = loadTestData();
  Matrix y_test = loadTestLabels();
  
  Matrix h1 = layer1.forward(X_test);
  Matrix a1 = activation1.forward(h1);
  Matrix h2 = layer2.forward(a1);
  Matrix y_pred = activation2.forward(h2);
  
  float test_loss = BinaryCrossEntropy::loss(y_pred, y_test);
  std::cout << "Test Loss: " << test_loss << std::endl;
  
  return 0;
}

6. 结论

通过使用 CUDA,我们能够在 GPU 上实现一个完整的神经网络,从而大大加速训练过程。这种低级别的实现不仅能够提高性能,还能帮助我们更深入地理解神经网络的工作原理。

虽然在实际应用中,我们通常会使用更高级的框架如 PyTorch,但是了解底层实现对于优化和调试复杂模型非常有帮助。此外,在某些情况下,自定义 CUDA 实现可以带来显著的性能提升,如 Flash Attention 所示。

参考文献

NVIDIA. (2024). CUDA C++ Programming Guide.
Harris, M. (2023). An Even Easier Introduction to CUDA. NVIDIA Developer Blog.
Paszke, A., et al. (2019). PyTorch: An Imperative Style, High-Performance Deep Learning Library. Advances in Neural Information Processing Systems.
Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems.

步子哥