在 Rust 中从零开始构建张量 (Part 1.2):视图操作
在上一部分,我们构建了张量库的基础。在这一部分中,我们将实现视图操作,如置换、重塑和切片。
视图操作是改变我们解释张量数据的方式,而无需复制或移动底层内存的转换。当你在 PyTorch 或 Candle 等框架中转置矩阵或切片张量时,你通常是在使用视图。原始数据保持不变,但索引逻辑会改变。
步长
回顾第一部分,步长指定了沿每个张量维度前进一步时在线性内存中跳跃的距离。
让我们通过一个具体的例子重新审视步长,以便更好地理解它们的工作原理。
对于形状为 [10, 9, 5, 13]
的张量
- 初始化
strides
为[1, 1, 1, 1]
- 从
i = 2
开始:strides[2] = strides[3] * shape[3] = 1 * 13 = 13
步长 = [1, 1, 13, 1]
i = 1
:strides[1] = strides[2] * shape[2] = 13 * 5 = 65
步长 = [1, 65, 13, 1]
- 最后,
i = 0
:strides[0] = strides[1] * shape[1] = 65 * 9 = 585
- 最终结果:
strides = [585, 65, 13, 1]
对于索引向量 [i, j, k, l]
i
在线性内存中跳跃9*5*13 = 585
个索引j
在线性内存中跳跃5*13 = 65
个索引k
在线性内存中跳跃13
个索引l
在线性内存中跳跃1
个索引
这是一个 3D 张量的动画
我们需要升级我们的 TensorShape 以缓存步长,这既是为了第一部分中提到的效率原因,也是为了实现我们将要实现的视图操作。
#[derive(Debug, Clone, PartialEq)]
struct TensorShape {
shape: Vec<usize>,
strides: Vec<usize>, // Changed
}
我们还需要一个新的函数来处理初始化
impl TensorShape {
fn new(shape: Vec<usize>) -> Self {
let mut strides = vec![1; shape.len()];
for i in (0..shape.len().saturating_sub(1)).rev() {
strides[i] = strides[i + 1] * shape[i + 1];
}
TensorShape { shape, strides }
}
// Rest of the impl ...
}
我们还需要更新 zeros
函数的实现
impl<T: Zero + Clone> Tensor<T> {
fn zeros(shape: Vec<usize>) -> Self {
let shape = TensorShape::new(shape); // Changed
let storage = TensorStorage::<T>::zeros(shape.size());
Tensor { shape, storage }
}
}
现在我们已经缓存了步长,我们可以更新 ravel_index
和 unravel_index
的实现来使用它们
impl TensorShape {
// Rest of the impl ...
fn ravel_index(&self, indices: &[usize]) -> usize {
if indices.len() != self.shape.len() {
panic!("Indices length does not match tensor shape dimensions.");
}
indices
.iter()
.zip(self.strides.iter()) // Changed
.map(|(&idx, &stride)| idx * stride) // Changed
.sum()
}
fn unravel_index(&self, index: usize) -> Vec<usize> {
if self.shape.is_empty() {
return vec![];
}
let mut indices = vec![0; self.shape.len()];
let mut remaining_index = index;
for (i, &stride) in self.strides.iter().enumerate() { // Changed
indices[i] = remaining_index / stride; // Changed
remaining_index %= stride; // Changed
}
indices
}
}
置换
随着步长处理的改进,我们现在可以实现第一个视图操作:置换。
我见过最常见的置换用例是在处理图像时。有两种常见的组织图像数据的方式
chw
: 通道(颜色),高度,宽度- 可以看作是
c
个独立的灰度图像堆叠在一起
- 可以看作是
hwc
: 高度,宽度,通道(颜色)- 可以看作是一个 2D 数组,其中每个元素包含
c
种颜色
- 可以看作是一个 2D 数组,其中每个元素包含
用 einops 符号表示,这种转换可以写成
h w c -> c h w
那么,我们如何高效地置换张量呢?
解决方案是同时置换 TensorShape
结构中的形状和步长元素。如果我们这样做,那么在计算索引时,我们会得到相同的线性索引。
这之所以有效,是因为索引和步长之间的点积
线性索引 = 索引 (点积) 步长 = 索引[0]*步长[0] + 索引[1]*步长[1] + 索引[2]*步长[2] + ...
当对两个向量应用相同的置换时,由于加法是可交换的,所以它保持不变
线性索引 = 索引 (点积) 步长 = 索引[1]*步长[1] + 索引[2]*步长[2] + ... + 索引[0]*步长[0] = 索引[0]*步长[0] + 索引[1]*步长[1] + 索引[2]*步长[2] + ...
我们可以将此功能添加到 TensorShape
中
impl TensorShape {
// Rest of the impl ...
fn permute(&self, permuted_indices: &[usize]) -> Self {
let shape = permuted_indices.iter().map(|&i| self.shape[i]).collect();
let strides = permuted_indices.iter().map(|&i| self.strides[i]).collect();
Self { shape, strides }
}
}
并将其添加到我们的 Tensor
实现中
impl<T: Clone> Tensor<T> {
fn permute(&self, permuted_indices: &[usize]) -> Self {
Tensor {
shape: self.shape.permute(permuted_indices),
storage: self.storage.clone(),
}
}
}
我们需要 T
实现 Clone
,否则我们无法克隆 TensorStorage
。请注意,我们正在克隆整个存储,尽管我们没有修改它。这种克隆是浪费的,因为底层数据没有改变,但避免它需要添加引用语义或智能视图,这是我们稍后将处理的问题。
在生产中:Candle
让我们看看 Candle 是如何实现 Permute 的。
以下是总结
pub fn permute(&self, idxs: &[usize]) -> Result<Self> {
let stride = self.stride();
let dims = self.shape().dims();
let mut perm_stride = stride.to_vec();
let mut perm_dims = dims.to_vec();
for (i, &idx) in idxs.iter().enumerate() {
perm_stride[i] = stride[idx];
perm_dims[i] = dims[idx];
}
Ok(Self {
shape: Shape::from(perm_dims),
stride: perm_stride,
start_offset: self.start_offset,
})
}
Candle 的做法与我们完全相同,只是它使用循环而不是 Rust 迭代器方法。
合并
对于合并操作,我们希望将多个相邻维度“合并”为一个维度。
本质上,我们选择将张量中的高维子张量线性化。
这在 MNIST 中出现,当二维图像需要转换为列向量时。
在 einops 符号中,合并批量灰度图像看起来像
b h w -> b (h w)
在底层,我们把相邻的维度大小相乘。
因此,如果我们有一个形状为 [3, 4, 5]
的张量,上述合并的输出形状将是 [3, 4*5] = [3, 20]
。
对于步长,我们通过仅保留合并的最后一个维度的步长来缩小维度数量,以保持正确的线性索引。
因此,对于我们上面的例子,我们要合并 1..=2
,步长 [20, 5, 1]
变为 [20, 1]
。我们只保留合并范围中最后一个维度(在本例中是维度 2)的步长。
现在我们理解了合并是如何在概念上运作的,以下是我们如何在 TensorShape
中实现它
use std::ops::RangeInclusive;
impl TensorShape {
// Rest of the impl ...
fn merge(&self, dim_range: RangeInclusive<usize>) -> Self {
let (start, end) = (*dim_range.start(), *dim_range.end());
assert!(
start <= end && end < self.shape.len(),
"Invalid dimension range for merge"
);
let merged_size = self.shape[dim_range.clone()].iter().product();
let merged_stride = self.strides[end];
let mut new_shape = Vec::with_capacity(self.shape.len() - (end - start));
let mut new_strides = Vec::with_capacity(self.strides.len() - (end - start));
new_shape.extend_from_slice(&self.shape[..start]);
new_shape.push(merged_size);
new_shape.extend_from_slice(&self.shape[end + 1..]);
new_strides.extend_from_slice(&self.strides[..start]);
new_strides.push(merged_stride);
new_strides.extend_from_slice(&self.strides[end + 1..]);
Self {
shape: new_shape,
strides: new_strides,
}
}
}
并将其添加到我们的 Tensor
实现中
impl<T: Clone> Tensor<T> {
// Rest of the impl ...
fn merge(&self, dim_range: RangeInclusive<usize>) -> Self {
Tensor {
shape: self.shape.merge(dim_range),
storage: self.storage.clone(),
}
}
}
这个 merge
函数接受一个包含范围。对于我们上面的例子,它看起来像 shape.merge(1..=2)
。
在生产中:Candle
在 Candle 中,他们没有调用 merge 函数,而是有一个名为 flatten 的函数。
这是核心逻辑
fn flatten_<D1: Dim, D2: Dim>(
&self,
start_dim: Option<D1>,
end_dim: Option<D2>,
) -> Result<Tensor> {
// Handle scalar tensors (0 dimensions) -> reshape to [1]
if self.rank() == 0 {
return self.reshape(1);
}
// Default to flattening all dimensions if not specified
let start_dim = start_dim.unwrap_or(0);
let end_dim = end_dim.unwrap_or(self.rank() - 1);
if start_dim < end_dim {
let dims = self.dims();
let mut new_dims = Vec::new();
// Keep dimensions before the flatten range
new_dims.extend(&dims[..start_dim]);
// Calculate product of flattened dimensions
let flattened_size: usize = dims[start_dim..=end_dim].iter().product();
new_dims.push(flattened_size);
// Keep dimensions after the flatten range
new_dims.extend(&dims[end_dim + 1..]);
self.reshape(new_dims)
} else {
// No flattening needed
Ok(self.clone())
}
}
这与我们上面所说的基本相同,但更简洁。
分割
对于分割操作,我们希望将一个维度“分割”成多个相邻的维度。
基本上,我们正在将 1D 子张量转换为多维子张量。
这在 MNIST 中出现,当我们 K 要将网络的输出转换回灰度图像时。
在 einops 符号中,分割批量灰度图像看起来像
b (h w) -> b h w
与合并不同,此函数需要有关输出子张量大小的额外信息。
在底层,我们将一个维度的大小分解成多个相邻的维度。
因此,如果我们有一个形状为 [3, 20]
的张量,上述使用 w=5
进行分割的输出形状将是 [3, 20 // 5, 5] = [3, 4, 5]
。
对于步长,我们插入与我们添加的维度数量相等的新维度。因此,对于我们上面的示例,步长 [20, 1]
变为 [20, 5, 1]
。我们基本上根据新形状重新计算相应的步长。
我们还将包括推断大小,因此在上面的示例中,h
将被推断,因为我们可以使用线性大小和其他变量来确定它应该是什么。我们最多可以有一个未知维度。这类似于 PyTorch 或 NumPy 等框架允许在重塑时推断一个维度的方式。
现在我们理解了分割是如何在概念上运作的,以下是我们如何在 TensorShape
中实现它
use std::collections::HashMap;
impl TensorShape {
// Rest of the impl ...
fn split(&self, dim: usize, shape: &[usize]) -> Self {
if dim >= self.shape.len() {
panic!("Dimension index out of bounds");
}
let original_size = self.shape[dim];
let original_stride = self.strides[dim];
// Calculate the product of non-zero sizes and find wildcard
let mut non_zero_product = 1usize;
let mut zero_index = None;
for (i, &size) in shape.iter().enumerate() {
if size == 0 {
if zero_index.is_some() {
panic!("Cannot have more than one wildcard (0) in split sizes");
}
zero_index = Some(i);
} else {
non_zero_product *= size;
}
}
// Create the final sizes, inferring wildcards
let mut final_sizes = shape.to_vec();
if let Some(zero_index) = zero_index {
if original_size % non_zero_product != 0 {
panic!(
"Cannot split dimension of size {} into sizes {:?}",
original_size, shape
);
}
let inferred_size = original_size / non_zero_product;
final_sizes[zero_index] = inferred_size;
}
let mut new_shape = Vec::new();
let mut new_strides = Vec::new();
// Add dimensions before the split
new_shape.extend_from_slice(&self.shape[..dim]);
new_strides.extend_from_slice(&self.strides[..dim]);
// Calculate strides for the split dimensions
let mut current_stride = original_stride;
for &size in final_sizes.iter().rev() {
new_strides.push(current_stride);
current_stride *= size;
}
// Reverse the strides we just added to maintain correct order
let start_idx = new_strides.len() - final_sizes.len();
new_strides[start_idx..].reverse();
// Add the split dimensions to shape
new_shape.extend_from_slice(&final_sizes);
// Add remaining dimensions after the split
if dim + 1 < self.shape.len() {
new_shape.extend_from_slice(&self.shape[dim + 1..]);
new_strides.extend_from_slice(&self.strides[dim + 1..]);
}
Self {
shape: new_shape,
strides: new_strides,
}
}
}
并将其添加到我们的 Tensor
实现中
impl<T: Clone> Tensor<T> {
// Rest of the impl ...
fn split(&self, dim: usize, shape: &[usize]) -> Self {
Tensor {
shape: self.shape.split(dim, shape),
storage: self.storage.clone(),
}
}
}
这个 split
函数接受一个维度索引和一个表示将该维度重塑成什么形状的列表,其中 0 表示推断维度。
所以对于我们上面的例子
shape.split(1, &[0, 5]);
重塑是合并和分割的组合
我们已经实现了合并和分割,但分割在大多数框架中并不常见。通常框架使用合并和重塑,但事实证明,每个重塑都可以分解为最多一个合并和一个分割操作的序列。
使用合并和分割并不比直接实现重塑更高效。但是,它在概念上更清晰,因为实践中发生的大多数重塑都是语义上有意义的分组和取消分组。而任意重塑,如 `[6, 8] -> [4, 12]`,不常见,因为它混合了通常不相关的元素。
分析重塑时,会出现四种不同的模式。
情况 1:纯分组(仅合并)
这发生在重塑只能通过我们的合并操作完成时。
例如,将 [2, 3, 4, 5]
重塑为 [6, 20]
,只能通过合并完成,因为它将前两个维度和后两个维度分组。
我们可以通过检查输出形状的“前缀积”(维度大小的累积积)是否包含在输入形状的前缀积中来检测这种情况。
例如:
- 输入 `[2, 3, 4, 5]` 的前缀积为:`[2, 2*3, 2*3*4, 2*3*4*5] = [2, 6, 24, 120]`
- 输出 `[6, 20]` 的前缀积为:`[6, 6*20] = [6, 120]`
请注意,输出中的每个元素,无论是 6
还是 120
,都出现在输入的前缀乘积中。
这告诉我们重塑可以通过以下方式实现:
- 合并维度 0-1:
[2, 3, 4, 5] -> [6, 4, 5]
- 合并维度 1-2:
[6, 4, 5] -> [6, 20]
情况 2:纯取消分组(仅分割)
这发生在重塑只能通过我们的分割操作完成时,与我们仅合并的重塑场景相反。
例如,将 [6, 20]
重塑为 [2, 3, 4, 5]
,只能通过分割完成,因为它分割了第一个维度和最后一个维度。
例如:
- 输入 `[6, 20]` 的前缀积为:`[6, 6*20] = [6, 120]`
- 输出 `[2, 3, 4, 5]` 的前缀积为:`[2, 2*3, 2*3*4, 2*3*4*5] = [2, 6, 24, 120]`
请注意,输入中的每个元素,无论是 6
还是 120
,都出现在输出的前缀乘积中。
这告诉我们重塑可以通过以下方式实现:
- 分割维度 0:
[6, 20] -> [2, 3, 20]
- 分割维度 2:
[2, 3, 20] -> [2, 3, 4, 5]
情况 3:先取消分组后分组(先分割后合并)
这是第一个更复杂的情况,需要两种操作。也是最后一个语义上有意义的重塑方式。
让我们考虑将 [6, 8]
重塑为 [2, 36]
。
首先,我们将 [6, 8]
分解为两组素因子:[(2, 3), (2, 2, 2)]
。
然后我们将 [2, 36]
分解为两组素因子:[(2), (2, 2, 2, 3)]
。
每个正整数(大于一)都可以被分解,并由其素因子唯一描述。
例如:
6 = 2 * 3
8 = 2 * 2 * 2
125 = 5 * 5 * 5
120 = 2 * 2 * 2 * 3 * 5
这些集合列表可以被展平为多种不同的排列方式。
例如:
[6, 8]
可以变成
[2, 3, 2, 2, 2]
[3, 2, 2, 2, 2]
[2, 36]
可以变成
[2, 2, 2, 2, 3]
[2, 2, 2, 3, 2]
[2, 2, 3, 2, 2]
[2, 3, 2, 2, 2]
.
如果这两组之间有一个共同的扁平列表,我们就可以解决重塑问题。
计算上有趣的是,以这种方式重新分配素因子的正确方法计算成本惊人地高,这意味着我们必须检查输入和输出列表中每个元素的素因子的每个排列。
实际的解决方案是使用随机搜索方法。我们尝试每组的随机排列,然后检查是否可以将它们分组,使展平列表匹配。对于大多数实际张量形状,这种随机方法可以很快找到解决方案,因为大多数张量形状包含许多小素因子的副本。
情况 4:先分组后取消分组(先合并后分割)
如果情况 3 失败,总会有一个有保障的备用策略。与大多数框架中的做法类似,只要张量大小相同,我们就会强制覆盖形状。
这可以被认为是将所有维度合并为一个平面维度,然后将该维度分割为所需的输出形状。
这总是有效的,但我们失去了维度分组和取消分组的语义意义。
例如,考虑将 [6, 8]
重塑为 [4, 12]
。6 和 8 的素因子无论如何排列都无法将 4 放在开头。
[6, 8] ~ [(2, 3), (2, 2, 2)]
但是如果我们合并这两个维度
[6, 8] -> [48]
然后我们可以分割成任何我们想要的形状,只要维度大小的乘积等于 48
[48] -> [4, 12]
虽然这种方法普遍适用,但更早的案例(纯分组、纯取消分组或先分割后合并)更受青睐,因为它们保留了语义意义。
切片
现在我们来讨论切片、窗口化或获取子张量。
你可能在框架中看到过这样的语法
T[:, :4, 2:10, :]
这很有趣,将引出 TensorShape
的最后一个组件。
让我们从起始范围开始。
例如:
T[:3, :5, :7, :9]
让我们看一个 2D 例子,一个形状为 [4, 4]
的矩阵
0 1 2 3
4 5 6 7
8 9 a b
c d e f
如果我们对第一个维度 [:2]
和第二个维度 [:3]
进行切片会怎样?
0 1 2
4 5 6
形状已改变,现在是 [2, 3]
。然而,步长保持不变,因为底层内存布局没有改变。只要索引保持在新形状的边界内,行之间仍然是 4
的跳跃,列之间仍然是 1
的跳跃。
但现在让我们引入从不同索引开始的情况。
回到原始矩阵
0 1 2 3
4 5 6 7
8 9 a b
c d e f
如果我们切片第一个维度 [2:]
和第二个维度 [3:]
呢?
我们会得到
b
f
现在的形状是 [2, 1]
,步长保持不变。但是如果我们尝试索引张量的 [0, 0]
,我们会得到错误的值(仍然是 0
)。然而,如果我们将子张量 [2, 3]
的左上角设置为偏移量,我们可以将其添加到任何索引:[2, 3] + [0, 0] = [2, 3]
,这将给我们 b
,而 [2, 3] + [1, 0] = [3, 3]
,这将给我们 f
。
但是将整个偏移量存储为 Vec
将会效率低下。如果我们能证明 ravel_index
函数是线性的,那么它将具有以下属性
ravel_index(nd_index_a + nd_index_b) = ravel_index(nd_index_a) + ravel_index(nd_index_b)
这意味着我们不必为偏移量存储整个 Vec
,而只需将线性偏移量缓存到我们的 1D 张量中。
这被证明是真的。以下是二维的快速证明概要
ravel_index(a + b) = (a + b)[0]*strides[0] + (a + b)[1]*strides[1]
ravel_index(a + b) = (a[0] + b[0])*strides[0] + (a[1] + b[1])*strides[1]
ravel_index(a + b) = a[0]*strides[0] + b[0]*strides[0] + a[1]*strides[1] + b[1]*strides[1]
ravel_index(a + b) = a[0]*strides[0] + a[1]*strides[1] + b[0]*strides[0] + b[1]*strides[1]
ravel_index(a + b) = ravel_index(a) + ravel_index(b)
更高维度的证明遵循相同的模式,留给读者作为练习。
让我们更新 TensorShape
以包含此偏移量
#[derive(Debug, Clone, PartialEq)]
struct TensorShape {
shape: Vec<usize>,
strides: Vec<usize>,
linear_offset: usize,
}
让我们创建一个切片函数
impl TensorShape {
// Rest of the impl ...
fn slice(&self, dim: usize, range: RangeInclusive<usize>) -> Self {
if dim >= self.shape.len() {
panic!("Dimension index out of bounds");
}
let start = *range.start();
let end = *range.end();
if start > end || end >= self.shape[dim] {
panic!("Invalid slice range for dimension {}", dim);
}
let mut new_shape = self.shape.clone();
new_shape[dim] = end - start + 1; // inclusive range
let additional_offset = start * self.strides[dim];
Self {
shape: new_shape,
strides: self.strides.clone(),
linear_offset: self.linear_offset + additional_offset,
}
}
}
并将其添加到我们的 Tensor
实现中
impl<T: Clone> Tensor<T> {
// Rest of the impl ...
fn slice(&self, dim: usize, range: RangeInclusive<usize>) -> Self {
Tensor {
shape: self.shape.slice(dim, range),
storage: self.storage.clone(),
}
}
}
由于我们为 TensorShape
添加了一个新字段,我们还需要修改 new
、permute
、merge
和 split
的返回值。例如
// ...
Self {
shape: new_shape,
strides: new_strides,
linear_offset: self.linear_offset, // Updated!
}
在生产中:Candle
Candle 在 Layout
结构体 中包含了此偏移量
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Layout {
shape: Shape,
stride: Vec<usize>,
start_offset: usize, // Here!
}
跳过
在上一节中,您可能注意到我遗漏了 Python 中切片的最后一部分,即您可以指定步长来跳过元素
T[::2, ::3, ::4]
这会取第一个维度中每隔一个元素,第二个维度中每隔两个元素,第三个维度中每隔三个元素。
我将其移入一个单独的函数,以使实现更简洁,因为它执行的操作与基本切片有根本不同。
当我们跳过元素时,该维度的形状会发生变化。例如,如果我们有 10 个元素并取每第三个元素(步长=3),结果中将有 ceil(10/3) = 4
个元素。
保持高效的关键见解是,要跳过元素,我们将相应的步长乘以步长大小。这之所以有效,是因为步长表示沿维度每一步在内存中跳跃的距离,通过将步长乘以我们的步长,我们只是在进行更大的跳跃。
例如,考虑一个形状为 [4, 5, 6]
且默认步长为 [30, 6, 1]
的张量。如果我们想跳过第二个维度中每隔一个元素(步长为 2),我们将第二个步长乘以 2:[30, 6*2, 1]
= [30, 12, 1]
。
现在,当我们在第二个维度中前进一步时,我们在线性内存中跳跃 12 个位置而不是 6 个,有效地跳过了每隔一个元素。
这是实现
impl TensorShape {
// Rest of the impl ...
fn skip(&self, dim: usize, step: usize) -> Self {
// perform the equivalent of slicing with no range, but a step
if dim >= self.shape.len() {
panic!("Dimension index out of bounds");
}
let mut new_strides = self.strides.clone();
new_strides[dim] = new_strides[dim] * step;
let mut new_shape = self.shape.clone();
new_shape[dim] = new_shape[dim].div_ceil(step);
Self {
shape: new_shape,
strides: new_strides,
linear_offset: self.linear_offset,
}
}
}
并将其添加到我们的 Tensor
实现中
impl<T: Clone> Tensor<T> {
// Rest of the impl ...
fn skip(&self, dim: usize, step: usize) -> Self {
Tensor {
shape: self.shape.skip(dim, step),
storage: self.storage.clone(),
}
}
}
代码
要运行此博客文章中的代码
首先,克隆存储库
git clone git@github.com:greenrazer/easytensor.git
然后切换到此部分的标记提交:
git checkout part-2
然后运行测试
cargo test
查看 `src/` 中的代码。
后续步骤
我们现在已经实现了强大的视图操作,如置换、重塑和切片,它们为我们提供了灵活、零拷贝的方式来查看我们的张量数据。这些是 Candle 和 PyTorch 等实际张量库中使用的基本构建块。
在下一部分中,我们将探讨如何实现数据操作,如映射和归约。
同时,尝试创建不同形状的张量,并尝试置换、合并和分割。可视化底层内存如何保持固定而形状和步长如何改变是建立直觉的好方法。你越了解视图,高级操作和惰性计算就越自然。