【训练营】小模型训练支持

CMakeLists.txt

@@ -16,6 +16,7 @@ set(CMAKE_CXX_EXTENSIONS OFF)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 add_compile_options(-Wunused -Wunused-function)
+add_compile_options(-D_USE_MATH_DEFINES)
 
 # ------------------------------------------------------------------------------
 # GoogleTest (submodule)

example/mnist/main.cc

@@ -17,6 +17,9 @@
 #include "example/mnist/dataset.h"
 #include "example/mnist/net.h"
 
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/device.h"
+
 DEFINE_string(dataset, "", "mnist dataset path");
 DEFINE_int32(bs, 64, "batch size");
 DEFINE_int32(num_epoch, 1, "num epochs");
@@ -36,6 +39,7 @@ constexpr char kDeviceCUDA[] = "cuda";
 DEFINE_validator(device,
                  [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; });
 
+
 int main(int argc, char *argv[]) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
@@ -47,14 +51,15 @@ int main(int argc, char *argv[]) {
     auto test_dataset = std::make_shared<MNISTDataset>(FLAGS_dataset, false);
     DataLoader test_dataloader(test_dataset, FLAGS_bs);
 
-    auto network = MNIST();
+    auto network = std::make_shared<MNIST>();
     Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
     Device cpu_device = Device();
-    network.To(device);
+    network->To(device);
+    network->To(cpu_device);
 
     auto loss_fn = nn::CrossEntropyLoss();
     loss_fn.To(device);
-    auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr);
+    auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr);
 
     for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) {
         int train_idx = 0;
@@ -66,12 +71,33 @@ int main(int argc, char *argv[]) {
             auto new_image = std::make_shared<Tensor>(image->To(device));
             auto new_label = std::make_shared<Tensor>(label->To(device));
 
-            auto outputs = network.Forward({new_image});
+            auto outputs = network->Forward({new_image});
             optimizer.ZeroGrad();
 
             auto loss = loss_fn.Forward({outputs[0], new_label});
             loss[0]->Backward();
 
+            // 临时加：检查梯度
+            auto params = network->Parameters();
+            float total_grad = 0.0f;
+            for (auto &param : params) {
+                if (param->grad()) {
+                    float *grad_ptr = static_cast<float *>(param->grad()->DataPtr());
+                    for (int i = 0; i < param->grad()->NumElements(); ++i) {
+                        total_grad += std::abs(grad_ptr[i]);
+                    }
+                }
+            }
+            LOG(ERROR) << "total grad sum: " << total_grad;
+
+            // 临时加：检查参数更新前的值
+            float *first_param = static_cast<float *>(params[0]->DataPtr());
+            float before = first_param[0];
+
+            // 临时加：检查参数更新后的值
+            float after = first_param[0];
+            LOG(ERROR) << "param before: " << before << ", after: " << after;
+
             // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA
             // between forward and backward.
             auto loss_cpu = loss[0]->To(cpu_device);
@@ -104,7 +130,7 @@ int main(int argc, char *argv[]) {
         auto new_label = std::make_shared<Tensor>(label->To(device));
 
         auto label_cpu = label->To(cpu_device);
-        auto outputs = network.Forward({new_image});
+        auto outputs = network->Forward({new_image});
         auto output_cpu = outputs[0]->To(cpu_device);
         auto loss = loss_fn.Forward({outputs[0], new_label});
         auto loss_cpu = loss[0]->To(cpu_device);

example/mnist/net.cc

@@ -8,6 +8,7 @@
 
 #include "infini_train/include/nn/modules/activations.h"
 #include "infini_train/include/nn/modules/container.h"
+#include "infini_train/include/nn/modules/conv2d.h"
 #include "infini_train/include/nn/modules/linear.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/tensor.h"
@@ -16,16 +17,21 @@ namespace nn = infini_train::nn;
 
 MNIST::MNIST() {
     std::vector<std::shared_ptr<nn::Module>> layers;
-    layers.push_back(std::make_shared<nn::Linear>(784, 30));
+    layers.push_back(std::make_shared<nn::Conv2d>(1, 16, 3));
+    layers.push_back(std::make_shared<nn::Sigmoid>());
+    layers.push_back(std::make_shared<nn::Conv2d>(16, 32, 3));
     layers.push_back(std::make_shared<nn::Sigmoid>());
     modules_["sequential"] = std::make_shared<nn::Sequential>(std::move(layers));
-    modules_["linear2"] = std::make_shared<nn::Linear>(30, 10);
+    modules_["linear2"] = std::make_shared<nn::Linear>(32 * 24 * 24, 10);
 }
 
 std::vector<std::shared_ptr<infini_train::Tensor>>
 MNIST::Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) {
     CHECK_EQ(x.size(), 1);
-    auto x1 = (*modules_["sequential"])(x);
-    auto x2 = (*modules_["linear2"])(x1);
-    return x2;
-}
+    auto reshaped = x[0]->View({x[0]->Dims()[0], 1, 28, 28});
+    std::vector<std::shared_ptr<infini_train::Tensor>> x_reshaped = {reshaped};
+    auto x1 = (*modules_["sequential"])(x_reshaped);
+    auto x2 = x1[0]->View({x1[0]->Dims()[0], 32 * 24 * 24})->Contiguous();
+    auto x3 = (*modules_["linear2"])({x2});
+    return x3;
+}

infini_train/include/autograd/conv2d.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/autograd/function.h"
+
+namespace infini_train {
+class Tensor;
+}
+
+namespace infini_train::autograd {
+
+class Conv2d : public Function {
+public:
+    static constexpr char kType[] = "Conv2dFunction";
+
+    Conv2d(int64_t stride, int64_t padding) : Function(kType), stride_(stride), padding_(padding) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    int64_t stride_ = 1;
+    int64_t padding_ = 0;
+    bool bias_ = false;
+    int64_t in_channels_ = 0;
+    int64_t out_channels_ = 0;
+    int64_t kernel_size_ = 0;
+    std::vector<int64_t> input_dims_;
+};
+} // namespace infini_train::autograd

infini_train/include/nn/modules/conv2d.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/device.h"
+#include "infini_train/include/nn/modules/module.h"
+
+namespace infini_train {
+class Tensor;
+class Device;
+}
+
+namespace infini_train::nn {
+class Conv2d : public CloneableModule<Conv2d> {
+public:
+    static constexpr char kType[] = "Conv2d";
+    static constexpr char kParamWeightName[] = "weight";
+    static constexpr char kParamBiasName[] = "bias";
+
+    Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size,
+           int64_t stride = 1, int64_t padding = 0, bool bias = true, Device device = Device());
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+
+private:
+    void ResetParameters();
+    int64_t in_channels_ = 0;
+    int64_t out_channels_ = 0;
+    int64_t kernel_size_ = 0;
+    int64_t stride_ = 1;
+    int64_t padding_ = 0;
+    bool bias_ = true;
+};
+} // namespace infini_train::nn

infini_train/include/utils/string_utils.h

@@ -2,6 +2,7 @@
 
 #include <string>
 #include <vector>
+#include <cstdint>
 
 namespace infini_train::utils {
 std::string DimsToString(const std::vector<int64_t> &dims);

infini_train/src/autograd/conv2d.cc

@@ -0,0 +1,84 @@
+#include "infini_train/include/autograd/conv2d.h"
+
+#include "glog/logging.h"
+
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::autograd {
+std::vector<std::shared_ptr<Tensor>> Conv2d::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_GE(input_tensors.size(), 2);
+    const auto &input = input_tensors[0];
+    const auto &weight = input_tensors[1];
+    std::shared_ptr<Tensor> bias = nullptr;
+    if (input_tensors.size() == 3) {
+        bias = input_tensors[2];
+    }
+
+    auto device = input->GetDevice().type();
+    return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(
+        {device, "Conv2dForward"}, input, weight, bias, stride_, padding_)};
+}
+
+void Conv2d::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                          const std::vector<std::shared_ptr<Tensor>> &) {
+    const auto &input = input_tensors[0];
+    const auto &weight = input_tensors[1];
+    bool need_input = ctx_.needs_input_grad().size() > 0 && ctx_.needs_input_grad()[0];
+    bool need_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1];
+
+    ctx_.SaveForBackward({need_weight ? input : nullptr, need_input ? weight : nullptr});
+
+    input_dims_ = input->Dims();
+    LOG(ERROR) << "SetupContext input_dims_ size: " << input_dims_.size();
+    in_channels_ = weight->Dims()[1];
+    out_channels_ = weight->Dims()[0];
+    kernel_size_ = weight->Dims()[2];
+    bias_ = input_tensors.size() == 3;
+}
+
+std::vector<std::shared_ptr<Tensor>> Conv2d::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    auto saved_tensors = ctx_.GetSavedTensors();
+    LOG(ERROR) << "Backward input_dims_ size: " << input_dims_.size();
+    CHECK_EQ(saved_tensors.size(), 2);
+    const auto &input = saved_tensors[0];
+    const auto &weight = saved_tensors[1];
+    LOG(ERROR) << "Backward weight ptr: " << weight.get();
+    if (weight) {
+        LOG(ERROR) << "Backward weight dims size: " << weight->Dims().size();
+    } else {
+        LOG(ERROR) << "Backward weight is null!";
+    }
+    const auto &grad_output = grad_outputs[0];
+
+    CHECK(!ctx_.needs_input_grad().empty()) << "needs_input_grad not populated in Conv2d::Backward";
+    bool need_grad_input = ctx_.needs_input_grad()[0];
+    bool need_grad_weight = ctx_.needs_input_grad().size() > 1 && ctx_.needs_input_grad()[1];
+    bool need_grad_bias = bias_ && ctx_.needs_input_grad().size() > 2 && ctx_.needs_input_grad()[2];
+
+    auto device = grad_output->GetDevice().type();
+
+    std::shared_ptr<Tensor> grad_input = nullptr;
+    std::shared_ptr<Tensor> grad_weight = nullptr;
+    std::shared_ptr<Tensor> grad_bias = nullptr;
+
+    if (need_grad_input) {
+        grad_input = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(
+            {device, "Conv2dBackwardInput"}, weight, grad_output, input_dims_, stride_, padding_);
+    }
+    if (need_grad_weight && weight) {
+    grad_weight = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(
+        {device, "Conv2dBackwardWeight"}, input, grad_output, weight->Dims(), stride_, padding_);
+    }
+    if (need_grad_bias) {
+        grad_bias = Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(
+            {device, "Conv2dBackwardBias"}, grad_output, out_channels_);
+    }
+
+    if (bias_) {
+        return {grad_input, grad_weight, grad_bias};
+    } else {
+        return {grad_input, grad_weight};
+    }
+}
+} // namespace infini_train::autograd