pytorch reshape view性能对比（以及einsum, matmul）

原创

已于 2022-07-09 21:44:05 修改 · 1.5k 阅读

标签

#pytorch #深度学习 #python

收录于

于 2022-07-07 16:05:55 首次发布

本文探讨了在PyTorch中reshape和view的区别，以及einsum和matmul在高维度矩阵运算时的性能表现。结论指出，reshape会改变数据在内存中的布局，而view仅改变索引；在有广播机制的情况下，einsum比matmul更快。同时，使用for循环实现复杂运算会大大降低效率，应避免使用。测试代码展示了各种操作的时间消耗。

reshape和view

reshape 和 view的具体功能和区别就不介绍了，可以直接查看官网，简单来说，就是reshape会改变各个维度在存储中的物理位置，而view的话，只会改变索引。

那么我们在使用的时候，到底是选哪个呢？

先说结论：差不多~

einsum和matmul

具体用法参见官网，当我们想在维度比较高的tensor上做复杂的矩阵乘法的时候，往往会选择用einsum，因为比较清晰简单，但是如果有办法用matmul的时候，是不是会犹豫两个的性能呢？

结论：在有matmul有broadcast的情况下，einsum更快。简单情况没有测试，但是简单情况直接用matmul比较方便。

如果用for循环来实现einsum可以实现的复杂功能的话，会慢很多，所以千万不要用for loop！！！！

代码

import torch
import time
from collections import defaultdict


def generate_tensor(shape, num=1):
    return [torch.rand(shape).cuda() for _ in range(num)]


def bench_view(datas, shape_to, repeat=1):
    out = []
    for _ in range(repeat):
        res = []
        for d in datas:
            tmp = d.view(shape_to)
            res.append(tmp)
        out.append(torch.cat(res, dim=0)[:2].shape)
    print(len(out