124 lines
4.7 KiB
Plaintext
124 lines
4.7 KiB
Plaintext
#include <iostream>
|
|
#include <vector>
|
|
#include <cmath>
|
|
#include <cassert>
|
|
|
|
// 辅助函数：打印张量形状信息
|
|
void print_shape(const std::vector<int>& shape) {
|
|
std::cout << "[";
|
|
for (size_t i = 0; i < shape.size(); ++i) {
|
|
std::cout << shape[i];
|
|
if (i != shape.size() - 1) std::cout << ", ";
|
|
}
|
|
std::cout << "]" << std::endl;
|
|
}
|
|
|
|
// 核心函数：将 3x896x896 图像转换为 4096x2x3x14x14 Patches
|
|
void image_to_patches(
|
|
const float* input_data, // 输入图像指针 [3, 896, 896]
|
|
float* output_data, // 输出 patches 指针 [4096, 2, 3, 14, 14]
|
|
int img_c, int img_h, int img_w, // 输入图像维度: 3, 896, 896
|
|
int patch_h, int patch_w, // Patch 尺寸: 14, 14
|
|
int target_n, int target_t // 目标维度: 4096, 2
|
|
) {
|
|
|
|
// 计算每行的 patch 数量
|
|
int patches_per_row = img_w / patch_w; // 896 / 14 = 64
|
|
int patches_per_col = img_h / patch_h; // 896 / 14 = 64
|
|
|
|
// 校验目标 N 是否匹配
|
|
assert(patches_per_row * patches_per_col == target_n && "计算出的 Patch 总数与目标的 N=4096 不符!");
|
|
|
|
// 遍历每一个 Patch
|
|
for (int p_idx = 0; p_idx < target_n; ++p_idx) {
|
|
// 计算当前 patch 在原图上的起始坐标
|
|
int row = p_idx / patches_per_row;
|
|
int col = p_idx % patches_per_row;
|
|
|
|
int start_y = row * patch_h;
|
|
int start_x = col * patch_w;
|
|
|
|
// 遍历时间/副本维度 T (0 到 1)
|
|
for (int t = 0; t < target_t; ++t) {
|
|
// 遍历 RGB 三个通道
|
|
for (int c = 0; c < img_c; ++c) {
|
|
// 遍历 Patch 内部像素
|
|
for (int ph = 0; ph < patch_h; ++ph) {
|
|
for (int pw = 0; pw < patch_w; ++pw) {
|
|
|
|
// --- 输入索引计算 [C, H, W] ---
|
|
int in_y = start_y + ph;
|
|
int in_x = start_x + pw;
|
|
// 防止越界 (虽然这里整除刚好不会越界)
|
|
if (in_y >= img_h || in_x >= img_w) continue;
|
|
|
|
size_t input_idx = (size_t)c * img_h * img_w +
|
|
(size_t)in_y * img_w +
|
|
(size_t)in_x;
|
|
|
|
// --- 输出索引计算 [N, T, C, h, w] ---
|
|
size_t output_idx = 0;
|
|
output_idx += (size_t)p_idx * target_t * img_c * patch_h * patch_w;
|
|
output_idx += (size_t)t * img_c * patch_h * patch_w;
|
|
output_idx += (size_t)c * patch_h * patch_w;
|
|
output_idx += (size_t)ph * patch_w;
|
|
output_idx += (size_t)pw;
|
|
|
|
// 赋值
|
|
output_data[output_idx] = input_data[input_idx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int main() {
|
|
// 1. 定义输入和输出形状
|
|
const int IMG_C = 3;
|
|
const int IMG_H = 896;
|
|
const int IMG_W = 896;
|
|
|
|
const int PATCH_H = 14;
|
|
const int PATCH_W = 14;
|
|
|
|
const int TARGET_N = 4096; // 64 * 64
|
|
const int TARGET_T = 2;
|
|
|
|
// 2. 分配内存
|
|
// 输入大小: 3 * 896 * 896 = 2,408,448
|
|
std::vector<float> input_image(IMG_C * IMG_H * IMG_W, 0.0f);
|
|
|
|
// 输出大小: 4096 * 2 * 3 * 14 * 14 = 4,825,344
|
|
// 注意：输出大小是输入的两倍，因为 T=2 维度相当于存了两份所有的 patches
|
|
std::vector<float> output_patches(TARGET_N * TARGET_T * IMG_C * PATCH_H * PATCH_W, 0.0f);
|
|
|
|
// 3. 填充一些测试数据 (例如给 R 通道赋值为 1.0)
|
|
for (int i = 0; i < IMG_H * IMG_W; ++i) {
|
|
input_image[i] = 1.0f; // R channel part (simplified indexing for init)
|
|
}
|
|
|
|
std::cout << "Input Image Shape: ";
|
|
print_shape({IMG_C, IMG_H, IMG_W});
|
|
std::cout << "Input Element Count: " << input_image.size() << std::endl;
|
|
|
|
// 4. 执行转换
|
|
image_to_patches(
|
|
input_image.data(),
|
|
output_patches.data(),
|
|
IMG_C, IMG_H, IMG_W,
|
|
PATCH_H, PATCH_W,
|
|
TARGET_N, TARGET_T
|
|
);
|
|
|
|
std::cout << "Output Patches Shape: ";
|
|
print_shape({TARGET_N, TARGET_T, IMG_C, PATCH_H, PATCH_W});
|
|
std::cout << "Output Element Count: " << output_patches.size() << std::endl;
|
|
|
|
// 5. 简单验证：检查第一个 patch 的第一个像素是否拷贝成功
|
|
// 根据我们的索引计算，output[0, 0, 0, 0, 0] 应该等于 input[0, 0, 0]
|
|
float test_val = output_patches[0];
|
|
std::cout << "Verification: output_patches[0] = " << test_val << " (Expected 1.0)" << std::endl;
|
|
|
|
return 0;
|
|
} |