diff --git a/CMakeLists.txt b/CMakeLists.txt index 709bc30f2..95bd6ca51 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -129,8 +129,15 @@ if(USE_CUDA) message(STATUS "Add USE_NCCL, use NCCL with CUDA") list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake) find_package(NCCL REQUIRED) - add_compile_definitions(USE_NCCL=1) - target_link_libraries(infini_train_cuda_kernels PUBLIC nccl) + # add_compile_definitions(USE_NCCL=1) + # target_link_libraries(infini_train_cuda_kernels PUBLIC nccl) + + # 不再用目录级 add_compile_definitions 影响所有target, 只标到真正需要NCCL的target + target_compile_definitions(infini_train_cuda_kernels PUBLIC USE_NCCL=1) + # 加头文件 + target_include_directories(infini_train_cuda_kernels PUBLIC ${NCCL_INCLUDE_DIRS}) + # 用FindNCCL找到完整库路径 + target_link_libraries(infini_train_cuda_kernels PUBLIC ${NCCL_LIBRARIES}) endif() endif() @@ -168,7 +175,13 @@ if(USE_CUDA) if(USE_NCCL) # If your core library code also directly references NCCL symbols (not only kernels), # keep this. Otherwise it's harmless. - target_link_libraries(infini_train PUBLIC nccl) + + # target_link_libraries(infini_train PUBLIC nccl) + + # nccl_impl.cc需要nccl.h和NCCL库 + target_compile_definitions(infini_train PUBLIC USE_NCCL=1) + target_include_directories(infini_train PUBLIC ${NCCL_INCLUDE_DIRS}) + target_link_libraries(infini_train PUBLIC ${NCCL_LIBRARIES}) endif() endif() diff --git a/example/mnist/dataset.cc b/example/mnist/dataset.cc index ee683f6d4..635e9c1e0 100644 --- a/example/mnist/dataset.cc +++ b/example/mnist/dataset.cc @@ -115,8 +115,16 @@ MNISTDataset::MNISTDataset(const std::string &dataset, bool train) std::pair, std::shared_ptr> MNISTDataset::operator[](size_t idx) const { CHECK_LT(idx, image_file_.dims[0]); - return {std::make_shared(image_file_.tensor, idx * image_size_in_bytes_, image_dims_), + // image_file_.tensor 在构造函数里已被转换为 float32， + // 每个样本的字节数必须按当前 dtype 计算（784 * sizeof(float) = 3136）， + // 不能用原始 uint8 文件的 784。 + const size_t image_sample_bytes = image_file_.tensor.SizeInBytes() / image_file_.dims[0]; + + // return {std::make_shared(image_file_.tensor, idx * image_size_in_bytes_, image_dims_), + // std::make_shared(label_file_.tensor, idx * label_size_in_bytes_, label_dims_)}; + return {std::make_shared(image_file_.tensor, idx * image_sample_bytes, image_dims_), std::make_shared(label_file_.tensor, idx * label_size_in_bytes_, label_dims_)}; + } size_t MNISTDataset::Size() const { return image_file_.dims[0]; } diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..276939c07 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -6,6 +6,13 @@ #include #include + +#include +#include +#include + + + #include "gflags/gflags.h" #include "glog/logging.h" @@ -17,12 +24,30 @@ #include "example/mnist/dataset.h" #include "example/mnist/net.h" + + +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/utils.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" + DEFINE_string(dataset, "", "mnist dataset path"); DEFINE_int32(bs, 64, "batch size"); DEFINE_int32(num_epoch, 1, "num epochs"); DEFINE_double(lr, 0.01, "learning rate"); DEFINE_string(device, "cpu", "device type (cpu/cuda)"); + + +DEFINE_string(model, "mlp", "model type (mlp/cnn)"); +DEFINE_string(optimizer, "sgd", "optimizer type (sgd/adam)"); + + +DEFINE_int32(nthread_per_process, 1, "Number of threads to use for each process. >1 enables DDP on visible CUDA devices."); + + using namespace infini_train; namespace { @@ -31,30 +56,181 @@ constexpr int kNumClasses = 10; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; + + + +constexpr char kModelMLP[] = "mlp"; +constexpr char kModelCNN[] = "cnn"; +constexpr char kOptimizerSGD[] = "sgd"; +constexpr char kOptimizerAdam[] = "adam"; + }; // namespace 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]); +DEFINE_validator(model, [](const char *, const std::string &value) { return value == kModelMLP || value == kModelCNN; }); +DEFINE_validator(optimizer, [](const char *, const std::string &value) { + return value == kOptimizerSGD || value == kOptimizerAdam; +}); + + +DEFINE_validator(nthread_per_process, [](const char *, int32_t value) { + return value >= 1; +}); + +// 重构： 把原来main()里面的训练逻辑整体搬进Train(rank), 入口结构对齐 example/gpt2/main.cc +void Train(const nn::parallel::Rank &rank) { + using namespace nn::parallel; + + // 设置thread-local global rank, DDP/ProcessGroup 依赖它定位当前线程归属 + global::thread_global_rank = rank.GlobalRank(); + + const int ddp_world_size = global::GetDataParallelSize(); + int ddp_rank = 0; + + Device device; + if (rank.IsParallel()) { + // 并行模式下忽略 --device, 按照 thread_rank选择CUDA卡 + device = Device(Device::DeviceType::kCUDA, global::GetDeviceIndex(rank.thread_rank())); + + auto *pg_factory = ProcessGroupFactory::Instance(device.type()); + if (ddp_world_size > 1) { + const auto *ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()), GetDataParallelGroupRanks(rank.GlobalRank())); + ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank()); + } + + } else { + // 保留单卡和CPU的选择 + device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); + } + const Device cpu_device = Device(); - auto train_dataset = std::make_shared(FLAGS_dataset, true); - DataLoader train_dataloader(train_dataset, FLAGS_bs); - // TODO(dcj): Add sampler & eval dataloader later. + auto train_dataset = std::make_shared(FLAGS_dataset, true); auto test_dataset = std::make_shared(FLAGS_dataset, false); - DataLoader test_dataloader(test_dataset, FLAGS_bs); - auto network = MNIST(); - Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); - Device cpu_device = Device(); - network.To(device); + // 分线程保障 + const size_t batch_size = static_cast(FLAGS_bs); + const size_t train_num_batches = (train_dataset->Size() + batch_size - 1) / batch_size; + if (ddp_world_size > 1) { + CHECK_EQ(train_num_batches % static_cast(ddp_world_size), 0) + << "MNIST train batches must be divisible by ddp_world_size. " + << "train_size=" << train_dataset->Size() << ", bs=" << FLAGS_bs + << ", batches=" << train_num_batches << ", ddp_world_size=" << ddp_world_size + << ". For 4 GPUs use --bs 60; for 2 GPUs --bs 64 is OK."; + } + + // DDP用DistributedDataLoader做rank分片；而单卡保持DataLoader + std::shared_ptr train_dataloader; + if (ddp_world_size > 1) { + train_dataloader = std::make_shared( + train_dataset, batch_size, static_cast(ddp_rank), static_cast(ddp_world_size) + ); + + } else { + train_dataloader = std::make_shared(train_dataset, batch_size); + } + + // 测试集仍然用DataLoader;评测在main rank做 + DataLoader test_dataloader(test_dataset, batch_size); + + //原模型选择逻辑 + std::shared_ptr network; + if (FLAGS_model == kModelMLP) { + network = std::make_shared(); + } else { + network = std::make_shared(); + } + network->To(device); + + + // 先To(device) 再DDP，最后optimizer + // DDP构造会基于当前参数注册backward hook; 如果之后再换参数/搬设备，hook有可能失效 + std::shared_ptr eval_module = network; + if (ddp_world_size > 1) { + DistributedDataParallelConfig ddp_config{.zero_stage = 0}; + // zero_stage=0: 纯DDP梯度AllReduce， 不用ZeRO + network = std::make_shared(network, rank, ddp_config); + + // 评测时 unwrap到内部module，避免eval forward触发DDP reducer的PrepareForBackward + eval_module = std::dynamic_pointer_cast(network)->module(); + CHECK(eval_module) << "Failed to unwrap DDP module for evaluation."; + } + + if (rank.IsMainRank()) { + LOG(ERROR) << "model: " << FLAGS_model << ", ddp_world_size: " << ddp_world_size + << ", ddp_rank: " << ddp_rank; + } + + auto loss_fn = std::make_shared(); + loss_fn->To(device); + + // 原有的optimizer选择逻辑；注意必须再DDP包装后创建，以保证拿到的是包装后的参数列表 + std::shared_ptr optimizer; + if (FLAGS_optimizer == kOptimizerSGD) { + optimizer = std::make_shared(network->Parameters(), FLAGS_lr); + } else { + optimizer = std::make_shared(network->Parameters(), FLAGS_lr); + } + + if (rank.IsMainRank()) { + LOG(ERROR) << "optimizer: " << FLAGS_optimizer << ", lr: " << FLAGS_lr; + } + + + // DDP： 1.只在main rank评测，避免多rank重复输出； 2.评测基于 eval_module (DDP时为内部 module), 训练仍走 network (DDP wrapper) + auto evaluate = [&](int epoch_id) { + if (!rank.IsMainRank()) { + return; + } + + //框架无 no_grad，评测前把参数临时换成Detach副本，避免 eval 建图污染依赖计数 + const auto named_params = eval_module->NamedParameters(); + for (const auto &[name, param] : named_params) { + const auto dot = name.rfind('.'); + *eval_module->mutable_module(name.substr(0, dot))->mutable_parameter(name.substr(dot + 1)) = param->Detach(); + } + + std::vector test_losses; + int correct = 0; + int total = 0; + for (const auto &[image, label] : test_dataloader) { + auto new_image = std::make_shared(image->To(device)); + auto new_label = std::make_shared(label->To(device)); + + + auto label_cpu = new_label->To(cpu_device); + auto outputs = eval_module->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); + + const int cur_batch_size = output_cpu.Dims()[0]; + for (int batch_idx = 0; batch_idx < cur_batch_size; ++batch_idx) { + auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; + const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; + const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; + + if (output_index == label_index) { + ++correct; + } + } + + total += cur_batch_size; + test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); + } + + // 在评测结束后恢复原始参数对象 + for (const auto &[name, param] : named_params) { + const auto dot = name.rfind('.'); + *eval_module->mutable_module(name.substr(0, dot))->mutable_parameter(name.substr(dot + 1)) = param; + } + + const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); + LOG(ERROR) << std::format("epoch {:2d} | test | Total: {}, Correct: {}, Accuracy: {:.4f}, AverageLoss: {:.6f}", epoch_id, total, correct, static_cast(correct) / total, avg_loss); + }; - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; @@ -62,71 +238,319 @@ int main(int argc, char *argv[]) { const auto epoch_start = std::chrono::high_resolution_clock::now(); - for (const auto &[image, label] : train_dataloader) { + for (const auto &[image, label] : *train_dataloader) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network.Forward({new_image}); - optimizer.ZeroGrad(); + auto outputs = network->Forward({new_image}); + optimizer->ZeroGrad(); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); - // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA - // between forward and backward. + // loss的D2H拷贝放在backward之后，避免前反向之间多一次同步 auto loss_cpu = loss[0]->To(cpu_device); - float current_loss = static_cast(loss_cpu.DataPtr())[0]; + const float current_loss = static_cast(loss_cpu.DataPtr())[0]; total_loss += current_loss; - if (train_idx % kNumItersOfOutputDuration == 0) { - LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() - << "] " - << " loss: " << current_loss; + + // 只有main rank打印训练日志; DDP下把进度换算成全局样本数 + if (rank.IsMainRank() && train_idx % kNumItersOfOutputDuration == 0) { + LOG(ERROR) << "epoch: " << epoch << ", [" + << train_idx * FLAGS_bs * static_cast(ddp_world_size) << "/" + << train_dataset->Size() << "] " + << " loss: " << current_loss; } - optimizer.Step(); + optimizer->Step(); train_idx += 1; } const auto epoch_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); - LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", - epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, - train_dataset->Size() / (duration_us / 1e6)); + // 只有main rank打印 epoch汇总； samples/s 按全局训练集大小/ rank0 epoch 时长估算 + if (rank.IsMainRank()) { + LOG(ERROR) << std::format( + "epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ddp {} | ({:.2f} ms | {:.0f} samples/s)", epoch, + FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, ddp_world_size, duration_us / 1e3f, + train_dataset->Size() / (duration_us / 1e6f) + ); + } + + evaluate(epoch); } +} - // TODO(dcj): Add no_grad() context manager later. - std::vector test_losses; - int correct = 0; - int total = 0; - for (const auto &[image, label] : test_dataloader) { - auto new_image = std::make_shared(image->To(device)); - auto new_label = std::make_shared(label->To(device)); - - auto label_cpu = label->To(cpu_device); - 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); - - const int batch_size = output_cpu.Dims()[0]; - for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; - const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; - const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; - if (output_index == label_index) { - ++correct; - } + +int main(int argc, char *argv[]) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + + // DDP仅支持CUDA； CPU多线程 DP 不在本任务范围内 + if (FLAGS_nthread_per_process > 1 && FLAGS_device != kDeviceCUDA) { + LOG(FATAL) << "DDP mode requires CUDA devices.Please use --device cuda with --nthread_per_process > 1."; + } + + // 对齐gpt2入口. MNIST只启用DP，而不用TP/SP/PP + nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, 1, false, 1, 1); + LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); + + // nthread_per_process>1时， 每个线程绑定一个rank, 执行一个Train(rank) + if (FLAGS_nthread_per_process > 1) { + std::vector threads; + for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), idx, nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); + threads.emplace_back(Train, rank); } - total += batch_size; - test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); + + for (auto &thread : threads) { thread.join(); } + } else { + nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(), FLAGS_nthread_per_process); + Train(rank); } - const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); - LOG(ERROR) << "Total: " << total << ", Correct: " << correct - << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; gflags::ShutDownCommandLineFlags(); google::ShutdownGoogleLogging(); return 0; } + + + + +// int main(int argc, char *argv[]) { +// gflags::ParseCommandLineFlags(&argc, &argv, true); +// google::InitGoogleLogging(argv[0]); + +// auto train_dataset = std::make_shared(FLAGS_dataset, true); +// DataLoader train_dataloader(train_dataset, FLAGS_bs); + +// // TODO(dcj): Add sampler & eval dataloader later. +// auto test_dataset = std::make_shared(FLAGS_dataset, false); +// DataLoader test_dataloader(test_dataset, FLAGS_bs); + +// // auto network = MNIST(); +// //////// auto network = std::make_shared(); + +// // 改动加入： mlp 和CNN +// std::shared_ptr network; +// if (FLAGS_model == kModelMLP) { +// network = std::make_shared(); +// } else { +// network = std::make_shared(); +// } +// LOG(ERROR) << "model: " << FLAGS_model; +// ///////////////////////////// + +// Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); +// Device cpu_device = Device(); +// // network.To(device); +// network->To(device); + +// // auto loss_fn = nn::CrossEntropyLoss(); +// // loss_fn.To(device); +// auto loss_fn = std::make_shared(); +// loss_fn->To(device); +// //auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); +// ////////////////auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); +// std::shared_ptr optimizer; +// if (FLAGS_optimizer == kOptimizerSGD) { +// optimizer = std::make_shared(network->Parameters(), FLAGS_lr); +// } else { +// optimizer = std::make_shared(network->Parameters(), FLAGS_lr); +// } +// LOG(ERROR) << "optimizer: " << FLAGS_optimizer << ", lr: " << FLAGS_lr; +// ////////////////////////////////// + + + + + +// // // ===== overfit single batch test ===== +// // { +// // auto first = *train_dataloader.begin(); +// // for (int step = 0; step < 300; ++step) { +// // auto new_image = std::make_shared(first.first->To(device)); +// // auto new_label = std::make_shared(first.second->To(device)); +// // auto outputs = network->Forward({new_image}); +// // optimizer.ZeroGrad(); +// // auto loss = loss_fn->Forward({outputs[0], new_label}); +// // loss[0]->Backward(); +// // auto loss_cpu = loss[0]->To(cpu_device); +// // if (step % 20 == 0) { +// // LOG(ERROR) << "[overfit] step " << step +// // << " loss: " << static_cast(loss_cpu.DataPtr())[0]; +// // } +// // optimizer.Step(); +// // } +// // } +// // // ======================================= +// // 增加:每个epoch结束后调用一次测试集评测 +// auto evaluate = [&](int epoch_id) { + +// const auto named_params = network->NamedParameters(); +// for (const auto &[name, param] : named_params) { +// const auto dot = name.rfind('.'); +// *network->mutable_module(name.substr(0, dot))->mutable_parameter(name.substr(dot + 1)) = param->Detach(); +// } + +// std::vector test_losses; +// int correct = 0; +// int total = 0; +// for (const auto &[image, label] : test_dataloader) { +// auto new_image = std::make_shared(image->To(device)); +// auto new_label = std::make_shared(label->To(device)); + +// auto label_cpu = label->To(cpu_device); +// 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); + +// const int batch_size = output_cpu.Dims()[0]; +// for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { +// auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; +// const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; +// const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; +// if (output_index == label_index) { +// ++correct; +// } +// } + +// total += batch_size; +// test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); +// } + + +// for (const auto &[name, param] : named_params) { +// const auto dot = name.rfind('.'); +// *network->mutable_module(name.substr(0, dot))->mutable_parameter(name.substr(dot + 1)) = param; +// } + +// const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); +// LOG(ERROR) << std::format("epoch {:2d} | test | Total: {}, Correct: {}, Accuracy: {:.4f}, AverageLoss: {:.6f}", epoch_id, total, correct, static_cast(correct) / total, avg_loss); +// }; +// ////////////////////////////// + + + +// for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { +// int train_idx = 0; +// float total_loss = 0.0; + +// const auto epoch_start = std::chrono::high_resolution_clock::now(); + +// for (const auto &[image, label] : train_dataloader) { +// auto new_image = std::make_shared(image->To(device)); +// auto new_label = std::make_shared(label->To(device)); + +// // // ===== debug pairing ===== +// // if (train_idx == 0 || train_idx == 1) { +// // auto lc = new_label->To(cpu_device); +// // auto ic = new_image->To(cpu_device); +// // const auto *lp = static_cast(lc.DataPtr()); +// // const float *ip = static_cast(ic.DataPtr()); +// // for (int i = 0; i < 16; ++i) { +// // double s = 0.0; +// // for (int j = 0; j < 784; ++j) { s += ip[i * 784 + j]; } +// // LOG(ERROR) << "[debug] batch " << train_idx << " sample " << i +// // << " | label = " << static_cast(lp[i]) +// // << " | image sum = " << s; +// // } +// // } +// // // =========================== + +// // auto outputs = network.Forward({new_image}); +// auto outputs = network->Forward({new_image}); +// //optimizer.ZeroGrad(); +// // 变成指针了 +// optimizer->ZeroGrad(); + +// // auto loss = loss_fn.Forward({outputs[0], new_label}); +// auto loss = loss_fn->Forward({outputs[0], new_label}); +// loss[0]->Backward(); + +// // // ===== debug start ===== +// // const auto params = network->Parameters(); // 关键：先把临时 vector 接住 +// // for (size_t pi = 0; pi < params.size(); ++pi) { +// // const auto &p = params[pi]; // 现在引用有效 +// // const auto g = p->grad(); +// // if (!g) { +// // LOG(ERROR) << "[debug] param " << pi << " grad is NULL"; +// // continue; +// // } +// // auto gc = g->To(cpu_device); +// // const float *gd = static_cast(gc.DataPtr()); +// // double sum = 0.0; +// // for (size_t j = 0; j < g->NumElements(); ++j) { sum += gd[j] * gd[j]; } +// // LOG(ERROR) << "[debug] param " << pi << " grad norm = " << std::sqrt(sum); +// // } +// // auto pb = network->Parameters()[0]->To(cpu_device); +// // LOG(ERROR) << "[debug] w0[0] before step: " << static_cast(pb.DataPtr())[0]; +// // // ===== debug end ===== + +// // 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); +// float current_loss = static_cast(loss_cpu.DataPtr())[0]; +// total_loss += current_loss; +// if (train_idx % kNumItersOfOutputDuration == 0) { +// LOG(ERROR) << "epoch: " << epoch << ", [" << train_idx * FLAGS_bs << "/" << train_dataset->Size() +// << "] " +// << " loss: " << current_loss; +// } + +// //optimizer.Step(); +// // 变成指针 +// optimizer->Step(); +// train_idx += 1; +// } + +// const auto epoch_end = std::chrono::high_resolution_clock::now(); +// const double duration_us = std::chrono::duration(epoch_end - epoch_start).count(); + +// LOG(ERROR) << std::format("epoch {:2d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} samples/s)", +// epoch, FLAGS_num_epoch - 1, total_loss / train_idx, FLAGS_lr, duration_us / 1e3f, +// train_dataset->Size() / (duration_us / 1e6)); + +// // 每个epoch结束评测一次 +// evaluate(epoch); +// } + +// // // TODO(dcj): Add no_grad() context manager later. +// // std::vector test_losses; +// // int correct = 0; +// // int total = 0; +// // for (const auto &[image, label] : test_dataloader) { +// // auto new_image = std::make_shared(image->To(device)); +// // auto new_label = std::make_shared(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 = loss_fn->Forward({outputs[0], new_label}); +// // auto loss_cpu = loss[0]->To(cpu_device); + +// // const int batch_size = output_cpu.Dims()[0]; +// // for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { +// // auto label_index = reinterpret_cast(label_cpu.DataPtr())[batch_idx]; +// // const auto *output_values = static_cast(output_cpu.DataPtr()) + batch_idx * kNumClasses; +// // const int output_index = std::max_element(output_values, output_values + kNumClasses) - output_values; +// // if (output_index == label_index) { +// // ++correct; +// // } +// // } +// // total += batch_size; +// // test_losses.push_back(static_cast(loss_cpu.DataPtr())[0]); +// // } +// // const auto avg_loss = std::accumulate(test_losses.begin(), test_losses.end(), 0.0) / test_losses.size(); +// // LOG(ERROR) << "Total: " << total << ", Correct: " << correct +// // << ", Accuracy: " << static_cast(correct) / total << ", AverageLoss: " << avg_loss; + +// gflags::ShutDownCommandLineFlags(); +// google::ShutdownGoogleLogging(); + +// return 0; +// } diff --git a/example/mnist/net.cc b/example/mnist/net.cc index 501fee7ef..88261bf4c 100644 --- a/example/mnist/net.cc +++ b/example/mnist/net.cc @@ -12,6 +12,9 @@ #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/tensor.h" +#include "infini_train/include/nn/modules/conv2d.h" + + namespace nn = infini_train::nn; MNIST::MNIST() { @@ -29,3 +32,34 @@ MNIST::Forward(const std::vector> &x) { auto x2 = (*modules_["linear2"])(x1); return x2; } + +// CNN手写数字识别网络 +MNISTCNN::MNISTCNN() { + modules_["conv1"] = std::make_shared(1, 16, 3); //in_channels=1, out_channels=16, kernel_size=3 + modules_["relu1"] = std::make_shared(); + modules_["conv2"] = std::make_shared(16, 32, 3); // in_channels=16, out_channels=32, kernel_size=3 + modules_["relu2"] = std::make_shared(); + + // 变成10个class + modules_["fc"]= std::make_shared(32 * 24 * 24, 10); +} + + +std::vector> +MNISTCNN::Forward(const std::vector> &x) { + CHECK_EQ(x.size(), 1); + const auto &input = x[0]; + const int64_t batch_size = input->Dims()[0]; + + // DataLoader: {bs, 28, 28} , Conv2d要求：4-D NCHW -> View {bs, 1, 28, 28} + // View通过Noop接入autograd，反向时梯度形状自动还原 + auto out = (*modules_["conv1"])({input->View({batch_size, 1, 28, 28})}); + out = (*modules_["relu1"])(out); + out = (*modules_["conv2"])(out); + out = (*modules_["relu2"])(out); + + // {bs, 32, 24, 24} ->拉平后接全连接层 + out = {out[0]->Flatten(1)}; + out = (*modules_["fc"])(out); + return out; +} diff --git a/example/mnist/net.h b/example/mnist/net.h index 5f4cfa33b..241842041 100644 --- a/example/mnist/net.h +++ b/example/mnist/net.h @@ -16,3 +16,13 @@ class MNIST : public infini_train::nn::Module { std::vector> Forward(const std::vector> &x) override; }; + + +// CNN手写数字识别网络 +class MNISTCNN : public infini_train::nn::Module { +public: + MNISTCNN(); + + std::vector> + Forward(const std::vector> &x) override; +}; diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..7a0bbac5b 100644 --- a/infini_train/include/autograd/activations.h +++ b/infini_train/include/autograd/activations.h @@ -21,4 +21,21 @@ class Sigmoid : public Function { const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; }; + + +class Relu : public Function { +public: + static constexpr char kType[] = "ReluFunction"; + + Relu() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + + + + } // namespace infini_train::autograd diff --git a/infini_train/include/autograd/conv2d.h b/infini_train/include/autograd/conv2d.h new file mode 100644 index 000000000..798f01ff6 --- /dev/null +++ b/infini_train/include/autograd/conv2d.h @@ -0,0 +1,37 @@ +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" + +namespace infini_train { + class Tensor; +} + +namespace infini_train::autograd { + +// Conv2d 的autograd函数节点 +// apply()会自动调用forward算输出，setupcontext存入反向所需的中间量 +// 反向传播框架会回调backward +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> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t stride_ = 1; + int64_t padding_ = 0; + bool bias_ = false; + + std::vector input_dims_; + std::vector weight_dims_; +}; + +} diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..85e11cc44 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -17,6 +17,18 @@ class Sigmoid : public CloneableModule { std::vector> Forward(const std::vector> &input_tensors) override; }; + + +class Relu : public CloneableModule { +public: + static constexpr char kType[] = "Relu"; + Relu() : CloneableModule(kType) {} + std::vector> Forward(const std::vector> &input_tensors) override; +}; + + + + class NewGELU : public CloneableModule { public: static constexpr char kType[] = "NewGELU"; diff --git a/infini_train/include/nn/modules/conv2d.h b/infini_train/include/nn/modules/conv2d.h new file mode 100644 index 000000000..5b7740e1b --- /dev/null +++ b/infini_train/include/nn/modules/conv2d.h @@ -0,0 +1,38 @@ +#pragma once + +#include +#include +#include + +#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 { +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> Forward(const std::vector> &input_tensors) override; + + bool has_bias() const { return bias_; } + +private: + void ResetParameters(); + + int64_t stride_ = 1; + int64_t padding_ = 0; + bool bias_ = true; +}; + + +} \ No newline at end of file diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5ea..cb0c5dafa 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -30,4 +30,34 @@ std::vector> Sigmoid::Backward(const std::vectorGetDevice().type(); return {Dispatcher::Instance().Call>({device, "SigmoidBackward"}, output, grad_output)}; } + + +std::vector> Relu::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReluForward"}, input)}; +} + +void Relu::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + // 存储前向输出的output, 反向用它判断哪些位置梯度为0 + const auto &output = output_tensors[0]; + ctx_.SaveForBackward({output}); +} + +std::vector> Relu::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + const auto &output = saved_tensors[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = output->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "ReluBackward"}, output, grad_output)}; +} + + + } // namespace infini_train::autograd diff --git a/infini_train/src/autograd/conv2d.cc b/infini_train/src/autograd/conv2d.cc new file mode 100644 index 000000000..c843d6aca --- /dev/null +++ b/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> Conv2d::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; + + // 按照输入设备分发到CPU和CUDA的Conv2dForward kernel + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "Conv2dForward"}, input, weight, bias, stride_, padding_)}; + +} + +void Conv2d::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + 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]; + + // 反向1grad_input需要weight, 反向2grad_weight需要input + // 哪边不需要梯度就存nullptr, 省显存/内存 + ctx_.SaveForBackward({need_weight ? input : nullptr, need_input ? weight : nullptr}); + + bias_ = input_tensors.size() == 3; + input_dims_ = input->Dims(); + weight_dims_ = weight->Dims(); +} + + +std::vector> Conv2d::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 2); + const auto &input = saved_tensors[0]; + const auto &weight = saved_tensors[1]; + + CHECK_EQ(grad_outputs.size(), 1); + 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 grad_input = nullptr; + std::shared_ptr grad_weight = nullptr; + std::shared_ptr grad_bias = nullptr; + + // 反向1： 输入的梯度 + if (need_grad_input) { + grad_input = Dispatcher::Instance().Call>({device, "Conv2dBackwardInput"}, weight, grad_output, input_dims_, stride_, padding_); + } + + // 反向2：权重的梯度 + if (need_grad_weight) { + grad_weight = Dispatcher::Instance().Call>({device, "Conv2dBackwardWeight"}, input, grad_output, weight_dims_, stride_, padding_); + } + + // 反向3：偏置的梯度 + if (need_grad_bias) { + grad_bias = Dispatcher::Instance().Call>({device, "Conv2dBackwardBias"}, grad_output, weight_dims_[0]); + } + + + if (bias_) { + return {grad_input, grad_weight, grad_bias}; + } else { + return {grad_input, grad_weight}; + } +} + + + +} \ No newline at end of file diff --git a/infini_train/src/kernels/cpu/conv2d.cc b/infini_train/src/kernels/cpu/conv2d.cc new file mode 100644 index 000000000..6e5d71e25 --- /dev/null +++ b/infini_train/src/kernels/cpu/conv2d.cc @@ -0,0 +1,251 @@ +// 每个函数块都只读取输入的tensor数据，算结果后，返回新的tensor +// 不在这里写autograd逻辑 +// 用autograd层调用即可 + +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { + // 公式与PyTorch一致: OH = (H + 2*padding - KH) / stride + 1 + inline int64_t ConvOutputSize(int64_t size, int64_t kernel, int64_t stride, int64_t padding) { + return (size + 2 * padding - kernel) / stride + 1; + } +} + +// namespace + +// Conv2dForward 前向 +// Conv2dBackwardInput 反向 +// Conv2dBackwardWeight 反向 +// Conv2dBackwardBias 反向 + +// 前向 +std::shared_ptr Conv2dForward(const std::shared_ptr &input, + const std::shared_ptr &weight, + const std::shared_ptr &bias, + int64_t stride, int64_t padding) { + const auto &input_dims = input->Dims(); + const auto &weight_dims = weight->Dims(); + CHECK_EQ(input_dims.size(), 4) << "Conv2d input must be 4-D (NCHW)"; + CHECK_EQ(weight_dims.size(), 4) << "Conv2d weight must be 4-D"; + CHECK_EQ(input_dims[1], weight_dims[1]) << "input channels must match weight channels"; + + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight_dims[0], kh = weight_dims[2], kw = weight_dims[3]; + // oh ow： 输出高度， 输出宽度 + const int64_t oh = ConvOutputSize(h, kh, stride, padding); + const int64_t ow = ConvOutputSize(w, kw, stride, padding); + CHECK_GT(oh, 0) << "Conv2d output height must be positive, check kernel/stride/padding"; + CHECK_GT(ow, 0) << "Conv2d output width must be positive, check kernel/stride/padding"; + + + auto output = std::make_shared(std::vector{n, o, oh, ow}, DataType::kFLOAT32); + output->Fill(0.0f); + + // NCHW布局 + const float *input_ptr = static_cast(input->DataPtr()); + const float *weight_ptr = static_cast(weight->DataPtr()); + const float *bias_ptr = bias ? static_cast(bias->DataPtr()) : nullptr; + float *output_ptr = static_cast(output->DataPtr()); + + // B->输出通道->输出高度->输出宽度 + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t oi = 0; oi < o; ++oi) { + for (int64_t oy = 0; oy < oh; ++oy) { + for (int64_t ox = 0; ox < ow; ++ox) { + // 累加器 + float acc = bias_ptr ? bias_ptr[oi] : 0.0f; + + // 输入通道->卷积核高->卷积核宽 + for (int64_t ci = 0; ci < c; ++ci) { + for (int64_t ky = 0; ky < kh; ++ky) { + // 输出坐标 (oy, ox)和卷积核位置(ky, kx) 对应输入坐标(iy, ix) + // iy = oy * stride + ky - padding + const int64_t iy = oy * stride + ky - padding; + if (iy < 0 || iy >= h) { + continue; + } + + for (int64_t kx = 0; kx < kw; ++kx) { + const int64_t ix = ox * stride + kx - padding; + if (ix < 0 || ix >= w) { + continue; + } + // input[ni, ci, iy, ix] * weight[oi, ci, ky, kx] + acc += input_ptr[( (ni*c + ci) * h + iy) * w + ix] * weight_ptr[( (oi*c + ci) * kh + ky) * kw + kx]; + } + } + } + + output_ptr[( (ni*o + oi) * oh + oy) * ow + ox] = acc; + + } + } + } + } + + return output; + +} + +// 反向1： 输入的梯度 +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims, + int64_t stride, int64_t padding) { + const auto &grad_dims = grad_output->Dims(); + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight->Dims()[0], kh = weight->Dims()[2], kw = weight->Dims()[3]; + const int64_t oh = grad_dims[2], ow = grad_dims[3]; + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + grad_input->Fill(0.0f); + + const float *weight_ptr = static_cast(weight->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t ci = 0; ci < c; ++ci) { + for (int64_t iy = 0; iy < h; ++iy) { + for (int64_t ix = 0; ix < w; ++ix) { + float acc = 0.0f; + for (int64_t oi = 0; oi < o; ++oi) { + for (int64_t ky = 0; ky < kh; ++ky) { + // 反解出：oy * stride = iy + padding - ky + const int64_t rem_y = iy + padding - ky; + // 如果除不尽说明(iy, ix)对卷积核这行没有贡献 + + if (rem_y < 0 || rem_y % stride != 0) { + continue; + } + const int64_t oy = rem_y / stride; + if (oy >= oh) { + continue; + } + + for (int64_t kx = 0; kx < kw; ++kx) { + const int64_t rem_x = ix + padding - kx; + if (rem_x < 0 || rem_x % stride != 0){ + continue; + } + const int64_t ox = rem_x / stride; + if (ox >= ow) { + continue; + } + acc += grad_output_ptr[( (ni*o + oi) * oh + oy) * ow + ox] * weight_ptr[( (oi*c + ci) * kh + ky) * kw + kx]; + } + + } + } + + grad_input_ptr[( (ni*c + ci) * h + iy) * w + ix] = acc; + } + } + } + } + return grad_input; +} + + +// 反向2： 权重的梯度 +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, + const std::vector &weight_dims, + int64_t stride, int64_t padding) { + const auto &input_dims = input->Dims(); + const auto &grad_dims = grad_output->Dims(); + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight_dims[0], kh = weight_dims[2], kw = weight_dims[3]; + const int64_t oh = grad_dims[2], ow = grad_dims[3]; + + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + grad_weight->Fill(0.0f); + + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_weight_ptr = static_cast(grad_weight->DataPtr()); + + for (int64_t oi = 0; oi < o; ++oi) { + for (int64_t ci = 0; ci < c; ++ci) { + for (int64_t ky = 0; ky < kh; ++ky) { + for (int64_t kx = 0; kx < kw; ++kx) { + float acc = 0.0f; + + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t oy = 0; oy < oh; ++oy) { + const int64_t iy = oy * stride + ky - padding; + if (iy < 0 || iy >= h) { + continue; + } + for (int64_t ox = 0; ox < ow; ++ox) { + const int64_t ix = ox * stride + kx - padding; + if (ix < 0 || ix >= w) { + continue; + } + acc += input_ptr[( (ni*c + ci) * h + iy) * w + ix] * grad_output_ptr[( (ni*o + oi) * oh + oy) * ow + ox]; + + } + } + } + + grad_weight_ptr[( (oi*c + ci) * kh + ky) * kw + kx] = acc; + } + } + } + } + return grad_weight; +} + + +// 反向3： 偏置的梯度 +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, + int64_t out_channels) { + const auto &grad_dims = grad_output->Dims(); + const int64_t n = grad_dims[0], oh = grad_dims[2], ow = grad_dims[3]; + + auto grad_bias = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32); + grad_bias->Fill(0.0f); + + + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_bias_ptr = static_cast(grad_bias->DataPtr()); + + for (int64_t oi = 0; oi < out_channels; ++oi) { + float acc = 0.0f; + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t oy = 0; oy < oh; ++oy) { + for (int64_t ox = 0; ox < ow; ++ox) { + acc += grad_output_ptr[( (ni * out_channels + oi) * oh + oy) * ow + ox]; + + } + } + } + grad_bias_ptr[oi] = acc; + } + return grad_bias; +} + + + +} +// 注册表dispatcher +// CPU版本跟CUDA版本名字相同, device不同，Dispatcher按照device分发 +#define REGISTER_CPU_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_CONV2D_KERNEL(Conv2dForward) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CPU_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CPU_CONV2D_KERNEL \ No newline at end of file diff --git a/infini_train/src/kernels/cpu/elementwise.cc b/infini_train/src/kernels/cpu/elementwise.cc index 213abdcfb..382312d9e 100644 --- a/infini_train/src/kernels/cpu/elementwise.cc +++ b/infini_train/src/kernels/cpu/elementwise.cc @@ -311,6 +311,21 @@ std::pair, std::shared_ptr> DivBackward(const st [](float x, float y) { return -x / (y * y); }); } +// ReLU前向: y = max(x, 0) +std::shared_ptr ReluForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return x > 0.0f ? x : 0.0f; }); +} + + +// ReLU反向：grad_input = grad_output * (y > 0 ? 1 : 0) +std::shared_ptr ReluBackward(const std::shared_ptr &output, const std::shared_ptr &grad_output) { + return UnaryBackward(grad_output, output, [](float y) { return y > 0.0f ? 1.0f : 0.0f; }); +} + + + + + } // namespace infini_train::kernels::cpu #define REGISTER_CPU_ELEMENTWISE_KERNEL(kernel_name) \ @@ -359,4 +374,7 @@ REGISTER_CPU_ELEMENTWISE_KERNEL(MulScalarBackward) REGISTER_CPU_ELEMENTWISE_KERNEL(DivForward) REGISTER_CPU_ELEMENTWISE_KERNEL(DivBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(ReluForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(ReluBackward) + #undef REGISTER_CPU_ELEMENTWISE_KERNEL diff --git a/infini_train/src/kernels/cuda/conv2d.cu b/infini_train/src/kernels/cuda/conv2d.cu new file mode 100644 index 000000000..a5d9bbb51 --- /dev/null +++ b/infini_train/src/kernels/cuda/conv2d.cu @@ -0,0 +1,308 @@ +// cpu版本用for循环 +// cuda版本用线程束 + +// Conv2dForward: 前向 +// Conv2dBackwardInput: 反向 +// Conv2dBackwardWeight: 反向 +// Conv2dBackwardBias: 反向 + +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +namespace { + // OH = (H + 2*padding - kH) / stride + 1 + inline int64_t ConvOutputSize(int64_t size, int64_t kernel, int64_t stride, int padding) { + return (size + 2 * padding - kernel) / stride + 1; + } + + // grid->block->thread, 让block固定为256个线程 + inline dim3 CalcGrid(int64_t numel) { + const int64_t block = 256; + return dim3(static_cast((numel + block - 1) / block)); + } + +} + +// 前向 +// __global__: 在GPU执行，在CPU启动; 每个线程算出张量的一个元素 out[ni, oi, oy, ox] +__global__ void Conv2dForwardKernel(const float *__restrict__ input, + const float *__restrict__ weight, + const float *__restrict__ bias, + float *__restrict__ output, + int64_t n, int64_t c, int64_t h, int64_t w, + int64_t o, int64_t kh, int64_t kw, + int64_t oh, int64_t ow, + int64_t stride, int64_t padding) { + // blockIdx.x当前线程所在的block的编号 + // blockDim.x 每个block的线程数（256） + // threadIdx.x 当前线程在block内的编号 + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + + // 越界判断 + if (idx >= n * o * oh * ow) { + return; + } + + // 反算坐标 + // idx = ((ni*o + oi) * oh + oy) * ow + ox + const int64_t ox = idx % ow; + const int64_t oy = (idx / ow) % oh; + const int64_t oi = (idx / (ow * oh)) % o; + const int64_t ni = idx / (ow * oh * o); + + float acc = bias ? bias[oi] : 0.0f; + + for (int64_t ci = 0; ci < c; ++ci) { + for (int64_t ky = 0; ky < kh; ++ky) { + const int64_t iy = oy * stride + ky - padding; + if (iy < 0 || iy >= h) { + continue; + } + + for (int64_t kx = 0; kx < kw; ++kx) { + const int64_t ix = ox * stride + kx - padding; + if (ix < 0 || ix >= w) { + continue; + } + + acc += input[((ni*c + ci) * h + iy) * w + ix] * weight[((oi*c + ci) * kh + ky) * kw + kx]; + } + } + } + + output[idx] = acc; +} + +// CPU端 +std::shared_ptr Conv2dForward(const std::shared_ptr &input, + const std::shared_ptr &weight, + const std::shared_ptr &bias, + int64_t stride, int64_t padding) { + const auto &input_dims = input->Dims(); + const auto &weight_dims = weight->Dims(); + CHECK_EQ(input_dims.size(), 4); + CHECK_EQ(weight_dims.size(), 4); + CHECK_EQ(input_dims[1], weight_dims[1]); + + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight_dims[0], kh = weight_dims[2], kw = weight_dims[3]; + const int64_t oh = ConvOutputSize(h, kh, stride, padding); + const int64_t ow = ConvOutputSize(w, kw, stride, padding); + CHECK_GT(oh, 0); + CHECK_GT(ow, 0); + + // CUDA的tensor构造需要指定device (I/O要在同一机器上) + auto output = std::make_shared(std::vector{n, o, oh, ow}, DataType::kFLOAT32, input->GetDevice()); + output->Fill(0.0f); + + const int64_t numel = n*o*oh*ow; + // 并行计算启动 + Conv2dForwardKernel<<>>( + static_cast(input->DataPtr()),// DataPtr()返回GPU显存指针 + static_cast(weight->DataPtr()), + bias ? static_cast(bias->DataPtr()) : nullptr, + static_cast(output->DataPtr()), n, c, h, w, o, kh, kw, oh, ow, stride, padding); + + // GPU kernel异步启动，所以同步要等待算完 + cudaDeviceSynchronize(); + return output; +} + + +// 反向1：输入的梯度 +__global__ void Conv2dBackwardInputKernel(const float *__restrict__ weight, + const float *__restrict__ grad_output, + float *__restrict__ grad_input, + int64_t n, int64_t c, int64_t h, int64_t w, + int64_t o, int64_t kh, int64_t kw, int64_t oh, int64_t ow, + int64_t stride, int64_t padding) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= n*c*h*w) { + return; + } + + const int64_t ix = idx % w; + const int64_t iy = (idx / w) % h; + const int64_t ci = (idx / (w * h)) % c; + const int64_t ni = idx / (w * h * c); + + float acc = 0.0f; + for (int64_t oi = 0; oi < o; ++oi) { + for (int64_t ky = 0; ky < kh; ++ky) { + const int64_t rem_y = iy + padding - ky; + + if (rem_y < 0 || rem_y % stride != 0) { + continue; + } + + const int64_t oy = rem_y / stride; + if (oy >= oh) { + continue; + } + + for (int64_t kx = 0; kx < kw; ++kx) { + const int64_t rem_x = ix + padding - kx; + if (rem_x < 0 || rem_x % stride != 0) { + continue; + } + + const int64_t ox = rem_x / stride; + if (ox >= ow) { + continue; + } + + acc += grad_output[((ni*o + oi) * oh + oy) * ow + ox] * weight[((oi * c + ci) * kh + ky) * kw + kx]; + } + } + } + grad_input[idx] = acc; +} + +std::shared_ptr Conv2dBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, + const std::vector &input_dims, + int64_t stride, int64_t padding) { + const auto &grad_dims = grad_output->Dims(); + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight->Dims()[0], kh = weight->Dims()[2], kw = weight->Dims()[3]; + const int64_t oh = grad_dims[2], ow = grad_dims[3]; + + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + + const int64_t numel = n*c*h*w; + Conv2dBackwardInputKernel<<>>( + static_cast(weight->DataPtr()), + static_cast(grad_output->DataPtr()), + static_cast(grad_input->DataPtr()), n, c, h, w, o, kh, kw, oh, ow, stride, padding + ); + + cudaDeviceSynchronize(); + return grad_input; +} + +// 反向2： 权重的梯度 +__global__ void Conv2dBackwardWeightKernel(const float *__restrict__ input, + const float *__restrict__ grad_output, + float *__restrict__ grad_weight, + int64_t n, int64_t c, int64_t h, int64_t w, + int64_t o, int64_t kh, int64_t kw, + int64_t oh, int64_t ow, + int64_t stride, int64_t padding){ + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= o*c*kh*kw) { + return; + } + + const int64_t kx = idx % kw; + const int64_t ky = (idx / kw) % kh; + const int64_t ci = (idx / (kw * kh)) % c; + const int64_t oi = idx / (kw * kh * c); + + float acc = 0.0f; + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t oy = 0; oy < oh; ++oy) { + const int64_t iy = oy * stride + ky - padding; + if (iy < 0 || iy >= h) { + continue; + } + + for (int64_t ox = 0; ox < ow; ++ox) { + const int64_t ix = ox * stride + kx - padding; + if (ix < 0 || ix >= w) { + continue; + } + + acc += input[((ni*c + ci) * h + iy) * w + ix] * grad_output[((ni * o + oi) * oh + oy) * ow + ox]; + } + } + } + grad_weight[idx] = acc; +} + + +std::shared_ptr Conv2dBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, + const std::vector &weight_dims, + int64_t stride, int64_t padding) { + const auto &input_dims = input->Dims(); + const auto &grad_dims = grad_output->Dims(); + const int64_t n = input_dims[0], c = input_dims[1], h = input_dims[2], w = input_dims[3]; + const int64_t o = weight_dims[0], kh = weight_dims[2], kw = weight_dims[3]; + const int64_t oh = grad_dims[2], ow = grad_dims[3]; + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_weight->Fill(0.0f); + + const int64_t numel = o*c*kh*kw; + Conv2dBackwardWeightKernel<<>>( + static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), + static_cast(grad_weight->DataPtr()), n, c, h, w, o, kh, kw, oh, ow, stride, padding + ); + + cudaDeviceSynchronize(); + return grad_weight; +} + + +// 反向3：偏置的梯度 +__global__ void Conv2dBackwardBiasKernel(const float *__restrict__ grad_output, + float *__restrict__ grad_bias, + int64_t n, int64_t o, int64_t oh, int64_t ow) { + const int64_t oi = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (oi >= o) { + return; + } + + float acc = 0.0f; + for (int64_t ni = 0; ni < n; ++ni) { + for (int64_t oy = 0; oy < oh; ++oy) { + for (int64_t ox = 0; ox < ow; ++ox) { + acc += grad_output[((ni*o + oi) * oh + oy) * ow + ox]; + } + } + } + grad_bias[oi] = acc; +} + +std::shared_ptr Conv2dBackwardBias(const std::shared_ptr &grad_output, + int64_t out_channels) { + const auto &grad_dims = grad_output->Dims(); + const int64_t n = grad_dims[0], oh = grad_dims[2], ow = grad_dims[3]; + + auto grad_bias = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, grad_output->GetDevice()); + grad_bias->Fill(0.0f); + + Conv2dBackwardBiasKernel<<>>( + static_cast(grad_output->DataPtr()), + static_cast(grad_bias->DataPtr()), n, out_channels, oh, ow + ); + cudaDeviceSynchronize(); + return grad_bias; + +} + + +} + +// 注册dispatcher +#define REGISTER_CUDA_CONV2D_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_CONV2D_KERNEL(Conv2dForward) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardInput) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardWeight) +REGISTER_CUDA_CONV2D_KERNEL(Conv2dBackwardBias) + +#undef REGISTER_CUDA_CONV2D_KERNEL \ No newline at end of file diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index 755f7433b..2913510a5 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1222,6 +1222,23 @@ std::shared_ptr SigmoidBackward(const std::shared_ptr &output, return UnaryBackward(grad_output, output, [] __device__(auto x) { return Mul(x, Sub(decltype(x){1}, x)); }); , INFINI_ALL_FLOATING_TYPES) } + + +// ReLU 前向： y = max(x, 0) +std::shared_ptr ReluForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return x > decltype(x){0} ? x : decltype(x){0}; }); , INFINI_ALL_FLOATING_TYPES) +} + + +// ReLU反向: grad_input = grad_output * (y > 0 ? 1 : 0) +std::shared_ptr ReluBackward(const std::shared_ptr &output, const std::shared_ptr &grad_output) { + DISPATCH(grad_output->Dtype(), return UnaryBackward(grad_output, output, [] __device__(auto y) { return y > decltype(y){0} ? decltype(y){1} : decltype(y){0}; }); , INFINI_ALL_FLOATING_TYPES) +} + + + + + } // namespace infini_train::kernels::cuda #define REGISTER_CUDA_ELEMENTWISE_KERNEL(kernel_name) \ @@ -1272,4 +1289,8 @@ REGISTER_CUDA_ELEMENTWISE_KERNEL(DivBackward) REGISTER_CUDA_ELEMENTWISE_KERNEL(SigmoidForward) REGISTER_CUDA_ELEMENTWISE_KERNEL(SigmoidBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReluForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReluBackward) + + #undef REGISTER_CUDA_ELEMENTWISE_KERNEL diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..b72238c5b 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -12,6 +12,14 @@ std::vector> Sigmoid::Forward(const std::vector()->Apply(input_tensors); } + +std::vector> Relu::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} + + + + std::vector> NewGELU::Forward(const std::vector> &x) { auto &input = x[0]; return {0.5 * input diff --git a/infini_train/src/nn/modules/conv2d.cc b/infini_train/src/nn/modules/conv2d.cc new file mode 100644 index 000000000..2f1e8f370 --- /dev/null +++ b/infini_train/src/nn/modules/conv2d.cc @@ -0,0 +1,48 @@ +#include "infini_train/include/nn/modules/conv2d.h" + +#include +#include +#include + + +#include "infini_train/include/autograd/conv2d.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + + +namespace infini_train::nn { +Conv2d::Conv2d(int64_t in_channels, int64_t out_channels, int64_t kernel_size, int64_t stride, int64_t padding, bool bias, Device device) : CloneableModule(kType), stride_(stride), padding_(padding), bias_(bias) { + device_ = device; + + // {0, C, kH, kW}与PyTorch一致 + parameters_[kParamWeightName] = std::make_shared(std::vector{out_channels, in_channels, kernel_size, kernel_size}, DataType::kFLOAT32, device_) -> RequiresGrad(); + + if (bias) { + parameters_[kParamBiasName] = std::make_shared(std::vector{out_channels}, DataType::kFLOAT32, device_) -> RequiresGrad(); + + } + ResetParameters(); +} + + +std::vector> Conv2d::Forward(const std::vector> &input_tensors) { + // Apply()自动建立计算图 + return std::make_shared(stride_, padding_) -> Apply(bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], parameters_[kParamBiasName]} : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); + +} + + +void Conv2d::ResetParameters() { + init::KaimingUniform(parameters_[kParamWeightName], sqrt(5.0f)); + if (bias_) { + const auto [fan_in, _] = init::CalculateFanInAndFanOut(parameters_[kParamWeightName]); + const float bound = fan_in > 0 ? 1.0 / sqrt(fan_in) : 0.0; + init::Uniform(parameters_[kParamBiasName], -bound, bound); + } +} + + + + +} \ No newline at end of file diff --git a/tests/autograd/test_autograd_conv2d.cc b/tests/autograd/test_autograd_conv2d.cc new file mode 100644 index 000000000..28d0ef97a --- /dev/null +++ b/tests/autograd/test_autograd_conv2d.cc @@ -0,0 +1,125 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/conv2d.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + + +using namespace infini_train; + +namespace { +// 在CPU上按值构造张量，再搬运到设备，最后打开requires_grad +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &values, Device device) { + auto cpu_tensor = std::make_shared(dims, DataType::kFLOAT32); + std::copy(values.begin(), values.end(), static_cast(cpu_tensor->DataPtr())); + + return std::make_shared(cpu_tensor->To(device))->RequiresGrad(); + +} + +} + + +class AutogradConv2dForwardTest : public infini_train::test::InfiniTrainTest{}; + +// 前向：1*1*3*3， 1个2*2卷积核, stride=1, padding=0, 带bias +// input = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] weight = [[1, 2], [3, 4]], bias = 1 +// out =[38， 48, 68, 78】 +TEST_P(AutogradConv2dForwardTest, Basic) { + auto input = MakeTensor({1, 1, 3, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9}, GetDevice()); + auto weight = MakeTensor({1, 1, 2, 2}, {1, 2, 3, 4}, GetDevice()); + auto bias = MakeTensor({1}, {1.0f}, GetDevice()); + + auto conv_fn = std::make_shared(1, 0); //stride = 1, padding = 0 + auto result = conv_fn->Apply({input, weight, bias}); + + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorNear(result[0], std::vector{38, 48, 68, 78}, 1e-5f); +} + + +TEST_P(AutogradConv2dForwardTest, MultiChannelWithPadding) { + auto input = std::make_shared(std::vector{1, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{1, 2, 3, 3}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + + auto conv_fn = std::make_shared(1, 1); // stride = 1, padding = 1 + auto result = conv_fn->Apply({input, weight}); + + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorNear(result[0], std::vector{8, 12, 8, 12, 18, 12, 8, 12, 8}, 1e-5f); +} + + +TEST_P(AutogradConv2dForwardTest, Stride2) { + auto input = std::make_shared(std::vector{1, 1, 4, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(1.0f); + auto weight = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + + auto conv_fn = std::make_shared(2, 0); // stride=2, padding = 0 + auto result = conv_fn->Apply({input, weight}); + + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorNear(result[0], 4.0f, 1e-5f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConv2dForwardTest); + +class AutogradConv2dBackwardTest : public infini_train::test::InfiniTrainTest {}; + +// 反向 +TEST_P(AutogradConv2dBackwardTest, WithBias) { + auto input = MakeTensor({1, 1, 3, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9}, GetDevice()); + auto weight = MakeTensor({1, 1, 2, 2}, {1, 2, 3, 4}, GetDevice()); + auto bias = MakeTensor({1}, {0.0f}, GetDevice()); + + auto conv_fn = std::make_shared(1, 0); //stride = 1, padding = 0 + auto result = conv_fn->Apply({input, weight, bias}); + EXPECT_EQ(result.size(), 1); + + auto grad_output = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grads = conv_fn->Backward({grad_output}); + + ASSERT_EQ(grads.size(), 3); + EXPECT_EQ(grads[0]->Dims(), (std::vector{1, 1, 3, 3})); + test::ExpectTensorNear(grads[0], std::vector{1, 3, 2, 4, 10, 6, 3, 7, 4}, 1e-5f); + EXPECT_EQ(grads[1]->Dims(), (std::vector{1, 1, 2, 2})); + test::ExpectTensorNear(grads[1], std::vector{12, 16, 24, 28}, 1e-5f); + + EXPECT_EQ(grads[2]->Dims(), (std::vector{1})); + test::ExpectTensorNear(grads[2], 4.0f, 1e-5f); + +} + + + +TEST_P(AutogradConv2dBackwardTest, NoBias) { + auto input = MakeTensor({1, 1, 3, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9}, GetDevice()); + auto weight = MakeTensor({1, 1, 2, 2}, {1, 2, 3, 4}, GetDevice()); + + auto conv_fn = std::make_shared(1, 0); // stride = 1, padding = 0 + auto result = conv_fn->Apply({input, weight}); + EXPECT_EQ(result.size(), 1); + + auto grad_output = std::make_shared(std::vector{1, 1, 2, 2}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grads = conv_fn->Backward({grad_output}); + + ASSERT_EQ(grads.size(), 2); + test::ExpectTensorNear(grads[0], std::vector{1, 3, 2, 4, 10, 6, 3, 7, 4}, 1e-5f); + test::ExpectTensorNear(grads[1], std::vector{12, 16, 24, 28}, 1e-5f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradConv2dBackwardTest); + diff --git a/tests/autograd/test_autograd_relu.cc b/tests/autograd/test_autograd_relu.cc new file mode 100644 index 000000000..9bcde3279 --- /dev/null +++ b/tests/autograd/test_autograd_relu.cc @@ -0,0 +1,61 @@ +// ReLU autograd节点的单元测试 （CPU/CUDA 双设备） +// Forward Backward各一个用例, 均为手算精确值 + +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/tensor.h" + +#include "tests/common/test_utils.h" + +using namespace infini_train; + +namespace { +// CPU上写入, 搬运到设备上， 打开requires_grad +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &values, Device device) { + auto cpu_tensor = std::make_shared(dims, DataType::kFLOAT32); + std::copy(values.begin(), values.end(), static_cast(cpu_tensor->DataPtr())); + + return std::make_shared(cpu_tensor->To(device))->RequiresGrad(); +} +} + +class AutogradReluTest : public infini_train::test::InfiniTrainTest {}; + +// 前向y = max(x, 0) +// 输入： {-1.5, 0.0, 0.5, 2.0}->{0.0, 0.0, 0.5, 2.0} + + +TEST_P(AutogradReluTest, Forward) { + auto input = MakeTensor({4}, {-1.5f, 0.0f, 0.5f, 2.0f}, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + + EXPECT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{4})); + test::ExpectTensorNear(result[0], std::vector{0.0f, 0.0f, 0.5f, 2.0f}, 1e-6f); + +} + + +TEST_P(AutogradReluTest, Backward) { + auto input = MakeTensor({4}, {-1.5f, 0.0f, 0.5f, 2.0f}, GetDevice()); + + auto relu_fn = std::make_shared(); + auto result = relu_fn->Apply({input}); + EXPECT_EQ(result.size(), 1); + + auto grad_output = std::make_shared(std::vector{4}, DataType::kFLOAT32, GetDevice(), true); + grad_output->Fill(1.0f); + auto grads = relu_fn->Backward({grad_output}); + + ASSERT_EQ(grads.size(), 1); + test::ExpectTensorNear(grads[0], std::vector{0.0f, 0.0f, 1.0f, 1.0f}, 1e-6f); +} + +INFINI_TRAIN_REGISTER_TEST(AutogradReluTest); \ No newline at end of file