c++ libtorch tensor 矩阵分块

代码示例

torch::Tensor a = torch::rand({ 2, 2, 4, 4 }, torch::kComplexDouble); // 2*2个4*4矩阵
torch::Tensor b11 = torch::rand({ 2, 2, 2, 2 }, torch::kComplexDouble);
torch::Tensor b12 = torch::rand({ 2, 2, 2, 2 }, torch::kComplexDouble);
torch::Tensor b21 = torch::rand({ 2, 2, 2, 2 }, torch::kComplexDouble);
torch::Tensor b22 = torch::rand({ 2, 2, 2, 2 }, torch::kComplexDouble);
std::cout << "a:" << a.sizes() << "\n" << a << std::endl;
std::cout << "b11:" << b11.sizes() << "\n" << b11 << std::endl;
std::cout << "b12:" << b12.sizes() << "\n" << b12 << std::endl;
std::cout << "b21:" << b21.sizes() << "\n" << b21 << std::endl;
std::cout << "b22:" << b22.sizes() << "\n" << b22 << std::endl;

a.narrow(2, 0, 2).narrow(3, 0, 2).copy_(b11); //top left a = a11 = b11
a.narrow(2, 0, 2).narrow(3, 2, 2).copy_(b12); //top right a = a12 = b12
a.narrow(2, 2, 2).narrow(3, 0, 2).copy_(b21); //bottom left a = a21 = b21
a.narrow(2, 2, 2).narrow(3, 2, 2).copy_(b22); //bottom right a = a22 = b22
std::cout << "a:" << a.sizes() << "\n" << a << std::endl;

narrow函数原型为

tensor.narrow(dim, start, length); 

• dim：我们想要缩小（narrow）的维度。
• start：在指定的维度上开始的位置。
• length：在指定的维度上要选择的元素数量。