torch、torch_sparse、dgl三者spmm显存占用对比

最新推荐文章于 2023-03-16 20:54:40 发布
最新推荐文章于 2023-03-16 20:54:40 发布
阅读量3.4k 收藏 5
点赞数 3
分类专栏： DGL & GNN
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/u013036495/article/details/115326606
版权

简单比较一下torch原生，torch_sparse(pyG依赖模块)以及DGL中spmm的显存占用和时间

版本：

torch 1.7
dgl 0.5
pyg 1.6

torch原生，显存占用主要是对结点敏感（bakward的时候用的稠密张量）。

# -*- coding: utf-8 -*-
"""
@Author: 
@Time: 2020/12/1
@Description: 
"""
import torch
import numpy as np
import torch_sparse
import dgl
import time

n = 60000
nnz = 50000

np.random.seed(123)
torch.manual_seed(123)

rows = np.random.randint(0, n, nnz)
cols = np.random.randint(0, n, nnz)

values = torch.randn(nnz).cuda().requires_grad_(True)

# torch.sparse.mm
X_sparse = torch.sparse_coo_tensor([rows, cols], values, size=(n, n)).cuda().requires_grad_(True)
Y_dense = torch.randn((n, 200)).cuda().requires_grad_(True)

print("memory allocated before multi: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before multi: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

# t = torch_sparse.spmm(torch.tensor([rows, cols], dtype=torch.long).cuda(), values, n, n, Y_dense).sum()
torch.cuda.synchronize()
start_time = time.time()
t = torch.sparse.mm(X_sparse, Y_dense).sum()

print("memory allocated before backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

t.backward()

print("t: {}".format(t))
print("values grad: {}".format(values.grad))
# print("x.grad: {0} y.grad: {1}".format(X_sparse.grad, Y_dense.grad))
print("memory allocated after backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated after backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

torch.cuda.synchronize()
print("spmm and backward time is {} s".format(time.time() - start_time))

n = 60000, nnz = 50000时

Using backend: pytorch
memory allocated before multi: 0.049234944 GB
max memory allocated before multi: 0.049234944 GB
memory allocated before backward: 0.049235456 GB
max memory allocated before backward: 0.147145216 GB
t: 1653.47900390625
values grad: tensor([ -9.5381, -11.1868,  -7.3291,  ...,  -7.0068,  10.7648, -11.1320],
       device='cuda:0')
memory allocated after backward: 0.097670144 GB
max memory allocated after backward: 14.546705408 GB
spmm and backward time is 0.5331666469573975 s

n = 61000, nnz = 50000时(说明显存对结点敏感)：

Using backend: pytorch
memory allocated before multi: 0.049800704 GB
max memory allocated before multi: 0.049800704 GB
memory allocated before backward: 0.049801216 GB
max memory allocated before backward: 0.148846592 GB
Traceback (most recent call last):
  File "/home/maqy/gnn/ginn_batch_compare/GINN-1130/memory_test.py", line 34, in <module>
    t.backward()
  File "/root/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 132, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 13.86 GiB (GPU 0; 14.76 GiB total capacity; 94.99 MiB already allocated; 13.72 GiB free; 100.00 MiB reserved in total by PyTorch)

Process finished with exit code 1

torch_sparse（pyG依赖的包），显存占用主要对边的数量敏感。

import torch
import numpy as np
import torch_sparse
import dgl
import time

n = 60000
nnz = 50000

np.random.seed(123)
torch.manual_seed(123)

rows = np.random.randint(0, n, nnz)
cols = np.random.randint(0, n, nnz)

values = torch.randn(nnz).cuda().requires_grad_(True)

# torch.sparse.mm
# X_sparse = torch.sparse_coo_tensor([rows, cols], values, size=(n, n)).cuda().requires_grad_(True)
Y_dense = torch.randn((n, 200)).cuda().requires_grad_(True)

print("memory allocated before multi: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before multi: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

torch.cuda.synchronize()
start_time = time.time()
t = torch_sparse.spmm(torch.tensor([rows, cols], dtype=torch.long).cuda(), values, n, n, Y_dense).sum()

print("memory allocated before backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

t.backward()

print("t: {}".format(t))
print("values grad: {}".format(values.grad))
# print("x.grad: {0} y.grad: {1}".format(X_sparse.grad, Y_dense.grad))
print("memory allocated after backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated after backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

torch.cuda.synchronize()
print("spmm and backward time is {} s".format(time.time() - start_time))

n = 60000, nnz = 50,000时（使用torch原生的会OOM）：

Using backend: pytorch
memory allocated before multi: 0.048434688 GB
max memory allocated before multi: 0.048434688 GB
memory allocated before backward: 0.089235456 GB
max memory allocated before backward: 0.17746944 GB
t: 1653.479248046875
values grad: tensor([ -9.5381, -11.1868,  -7.3291,  ...,  -7.0068,  10.7648, -11.1320],
       device='cuda:0')
memory allocated after backward: 0.096869888 GB
max memory allocated after backward: 0.20943616 GB
spmm and backward time is 0.03326892852783203 s

Process finished with exit code 0

但torch_sparse是使用的gather scatter模式，边的条数变多的话，显存会明显增长，比如：
n = 60000, nnz = 500,000时：

Using backend: pytorch
memory allocated before multi: 0.05023488 GB
max memory allocated before multi: 0.05023488 GB
memory allocated before backward: 0.458791424 GB
max memory allocated before backward: 0.90758144 GB
t: -14356.595703125
values grad: tensor([ 11.1778, -22.7249,  -4.5888,  ..., -24.6294,   4.4844,   4.5446],
       device='cuda:0')
memory allocated after backward: 0.100470272 GB
max memory allocated after backward: 1.662460416 GB
spmm and backward time is 0.17125725746154785 s

Process finished with exit code 0

n = 61000, nnz = 5,000,000时：

memory allocated before multi: 0.069206016 GB
max memory allocated before multi: 0.069206016 GB
memory allocated before backward: 4.149206528 GB
max memory allocated before backward: 8.197440512 GB
Traceback (most recent call last):
  File "/home/maqy/gnn/ginn_batch_compare/GINN-1130/memory_test.py", line 33, in <module>
    t.backward()
  File "/root/miniconda3/lib/python3.7/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 132, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 3.73 GiB (GPU 0; 14.76 GiB total capacity; 11.31 GiB already allocated; 2.50 GiB free; 11.32 GiB reserved in total by PyTorch)

Process finished with exit code 1

使用DGL的方式：

import torch
import numpy as np
import torch_sparse
import dgl
import time

n = 60000
nnz = 50000

np.random.seed(123)
torch.manual_seed(123)

rows = np.random.randint(0, n, nnz)
cols = np.random.randint(0, n, nnz)

values = torch.randn(nnz).cuda().requires_grad_(True)

# torch.sparse.mm
# X_sparse = torch.sparse_coo_tensor([rows, cols], values, size=(n, n)).cuda().requires_grad_(True)
Y_dense = torch.randn((n, 200)).cuda().requires_grad_(True)

# 注意，这里需要是cols在前，rows在后，计算出来的结果才是一样的
g = dgl.graph((cols, rows))
g = g.to("cuda:0")
g.srcdata['h'] = Y_dense
g.edata['e'] = values

print("memory allocated before multi: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before multi: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

torch.cuda.synchronize()
start_time = time.time()
t = dgl.ops.gspmm(g, 'mul', 'sum', lhs_data=g.srcdata['h'], rhs_data=g.edata['e']).sum()

print("memory allocated before backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

t.backward()

print("t: {}".format(t))
print("values grad: {}".format(values.grad))
# print("x.grad: {0} y.grad: {1}".format(X_sparse.grad, Y_dense.grad))
print("memory allocated after backward: {} GB".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated after backward: {} GB".format(torch.cuda.max_memory_allocated() / 10 ** 9))

torch.cuda.synchronize()
print("spmm and backward time is {} s".format(time.time() - start_time))

n = 60000, nnz = 50000时：

Using backend: pytorch
memory allocated before multi: 0.048434688 GB
max memory allocated before multi: 0.048434688 GB
memory allocated before backward: 0.0484352 GB
max memory allocated before backward: 0.09667072 GB
t: 1653.47900390625
values grad: tensor([ -9.5381, -11.1868,  -7.3291,  ...,  -7.0068,  10.7648, -11.1320],
       device='cuda:0')
memory allocated after backward: 0.096869888 GB
max memory allocated after backward: 0.145104896 GB
spmm and backward time is 0.02077507972717285 s

Process finished with exit code 0

n = 60000, nnz = 500,000时：

Using backend: pytorch
memory allocated before multi: 0.05023488 GB
max memory allocated before multi: 0.05023488 GB
memory allocated before backward: 0.050235392 GB
max memory allocated before backward: 0.098470912 GB
t: -14356.599609375
values grad: tensor([ 11.1778, -22.7249,  -4.5888,  ..., -24.6294,   4.4844,   4.5446],
       device='cuda:0')
memory allocated after backward: 0.100470272 GB
max memory allocated after backward: 0.14870528 GB
spmm and backward time is 0.10973095893859863 s

Process finished with exit code 0

n = 60000, nnz = 5,000,000时：

Using backend: pytorch
memory allocated before multi: 0.069206016 GB
max memory allocated before multi: 0.069206016 GB
memory allocated before backward: 0.069206528 GB
max memory allocated before backward: 0.117442048 GB
t: 7377.28125
values grad: tensor([ 18.3335, -23.7182,  -4.3458,  ...,  13.7045,   9.3865, -12.4589],
       device='cuda:0')
memory allocated after backward: 0.13744128 GB
max memory allocated after backward: 0.185676288 GB
spmm and backward time is 1.038252353668213 s

Process finished with exit code 0

用于验证DGL和torch_sparse结果一样的代码：

import torch
import numpy as np
import torch_sparse
import dgl
import dgl.function as fn

torch.manual_seed(123)
n = 10
nnz = 2

rows = np.random.randint(0, n, nnz)
cols = np.random.randint(0, n, nnz)

a = torch.randn(nnz)
values = a.cuda().requires_grad_(True)
values2 = a.cuda().requires_grad_(True)

# torch.sparse.mm
Y_dense = torch.randn((n, 2)).cuda().requires_grad_(True)

g = dgl.graph([])
g.add_nodes(n)
g.add_edges(cols, rows)
g = g.to("cuda:0")
g.srcdata['h'] = Y_dense
g.edata['e'] = values

print("memory allocated before multi: {}".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before multi: {}".format(torch.cuda.max_memory_allocated() / 10 ** 9))

t1 = dgl.ops.gspmm(g, 'mul', 'sum', lhs_data=g.srcdata['h'], rhs_data=g.edata['e']).sum()

t2 = torch_sparse.spmm(torch.tensor([rows, cols], dtype=torch.long).cuda(), values2, n, n, Y_dense).sum()

print("memory allocated before backward: {}".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated before backward: {}".format(torch.cuda.max_memory_allocated() / 10 ** 9))

t1.backward()
t2.backward()

print("t1: {}".format(t1))
print("t2: {}".format(t2))
print("values grad: {}".format(values.grad))
print("values2 grad: {}".format(values2.grad))
# print("x.grad: {0} y.grad: {1}".format(X_sparse.grad, Y_dense.grad))
print("memory allocated after backward: {}".format(torch.cuda.memory_allocated() / 10 ** 9))
print("max memory allocated after backward: {}".format(torch.cuda.max_memory_allocated() / 10 ** 9))

输出：

memory allocated before multi: 1.536e-06
max memory allocated before multi: 1.536e-06
memory allocated before backward: 3.584e-06
max memory allocated before backward: 4.096e-06
t1: 0.010851062834262848
t2: 0.010851062834262848
values grad: tensor([1.2198, 1.2198], device='cuda:0')
values2 grad: tensor([1.2198, 1.2198], device='cuda:0')
memory allocated after backward: 4.096e-06
max memory allocated after backward: 8.704e-06

Process finished with exit code 0

所需时间测试

速度上，n=60000,nnz=50000的情况下，计时spmm和backward的时间
torch原生：
spmm and backward time is 0.5331666469573975 s

torch_sparse:
spmm and backward time is 0.03326892852783203 s

DGL:
spmm and backward time is 0.02077507972717285 s

可以看到DGL的方式确实会快一些。

n = 60000, nnz = 500,000时：
torch原生：
spmm and backward time is 0.5481641292572021 s

torch_sparse:
spmm and backward time is 0.17125725746154785 s

DGL:
spmm and backward time is 0.10973095893859863 s

先云
关注
  • 3
    点赞
  • 踩
  • 5
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
专栏目录
CUDA编程——GPU架构，由sp，sm，thread，block，grid，warp说起
Zhang Junior 的博客
01-19 5万+
掌握部分硬件知识，有助于程序员编写更好的CUDA程序，提升CUDA程序性能，本文目的是理清sp，sm，thread，block，grid，warp之间的关系。由于作者能力有限，难免有疏漏，肯定读者批评指正。
稀疏矩阵 以及 torch_sparse 整理
qq_40926715的博客
09-03 993
稀疏矩阵处理
torch中的spmm
闲看庭前雪的博客
05-19 1万+
系列文章目录 本系列记录自己的代码学习知识 torch.matmul的前后两个矩阵维度不同的小结 torch中的transpose和view的不同 torch中的spmm系列文章目录前言一、torch中的spmm总结 前言 一、torch中的spmm 我自己运行的结果，spmm就是矩阵乘法 例： z = y.transpose(0,1) output = torch.spmm(y, z) print('z:',z,'y:',y) print('output:',output) z: tensor
安装torch-scatter/torch-sparse无法继续的解决方法
qq_43481435的博客
10-08 1万+
pip install无法安装torch-scatter/torch-sparse的解决方法
torch,torchvision,torch-sparse，torch-scatter版本问题
qq_41220883的博客
04-22 6165
最近安装torch，torch-sparse，torch-scatter，torch-cluster，torch-spline-conv，遇到版本不匹配的问题，试了很多，找到了解决办法。^_ ^，磕盐人呀，必须解决。 1.已有基础： Cuda:10.1 Cudnn:7.6.5 2.再此基础上安装：在官网https://pytorch.org/get-started/previous-versions/ 可以找到cuda对应的pytorch的版本，以及选择是否安装GPU。这里选择使用命令，#cuda=10
CUDA中grid、block、thread、warp与SM、SP的关系
ambition_zhou的博客
05-07 4199
首先概括一下这几个概念。其中SM（Streaming Multiprocessor）和SP（streaming Processor）是硬件层次的，其中一个SM可以包含多个SP。thread是一个线程，多个thread组成一个线程块block，多个block又组成一个线程网格grid。 现在就说一下一个kenerl函数是怎么执行的。一个kernel程式会有一个grid，grid底下又有数个block，每个block...
PyPI 官网下载 | torch_sparse-0.6.12.tar.gz
01-30
《PyPI官网下载 | torch_sparse-0.6.12.tar.gz：深度学习中的稀疏张量处理》 PyPI（Python Package Index）是Python开发者的重要资源库，它提供了丰富的Python库供用户下载和使用。在PyPI官网上，我们可以找到名为...
Python库 | torch_optim_sparse-0.1.1-py3-none-any.whl
03-25
`torch_optim_sparse-0.1.1-py3-none-any.whl` 是一个与Python相关的库文件，特别地，它是针对PyTorch框架的一个优化器扩展库。这个库的名称暗示了它可能专注于实现针对稀疏参数的优化算法，这对于处理大规模稀疏...
torch_sparse-0.6.2-cp37-cp37m-win_amd64whl.zip
12-26
1. **稀疏张量类（SparseTensor）**：这是Torch_Sparse的核心数据结构，它能有效地存储和处理非零元素较少的大矩阵，减少了内存占用和计算成本。 2. **矩阵运算**：包括加法、减法、乘法等基本运算，以及转置、切片...
torch_sparse-0.6.8-cp37-cp37m-win_amd64whl.zip
最新发布
12-29
此外，torch_sparse库还提供了一些高级功能，如`spmm()`（稀疏矩阵与密集矩阵乘法）和`spadd()`（稀疏张量的加法），这些操作在图的邻接矩阵计算和消息传递过程中起到了关键作用。结合PyTorch的强大动态计算图能力，...
torch_sparse-0.6.17+pt113cpu-cp310-cp310-linux_x86_64.whl.zip
12-22
`torch_sparse`提供了如`spadd`、`spmm`等操作，用于在GPU上高效地执行稀疏矩阵乘法和其他相关运算。例如： ```python import torch from torch_sparse import SparseTensor # 创建稀疏张量 row = torch.tensor([0...
傻瓜基础运算相关处理
小七
05-02 1858
说明 临时起意，记一下运算相关处理技巧 除法分母为零 numpy自动处理为inf.， 并给出runtime warning 这个在之前训练grarep遇到过多次，thunlp的OpenNE没给出处理方法，现在看spagan给出了，记一下 RuntimeWarning: divide by zero encountered in power SPAGAN处理： d_inv_sqrt = np.po...
torch.sparse稀疏矩阵
andy dennis的博客
02-11 3688
前言 这几天接触图神经网络，会发现有些图的边相对结点比较少，那么邻接矩阵很多元素都为0， 显然，我们可以优化存储策略，以节省内存和加快运算。本文udating… 例子 import torch # 邻接矩阵 adj_dense = torch.tensor([[2, 0], [0, 3]]) # 下标 [0, 0], [1, 1] indices = torch.tensor([[0, 1], [0, 1]]) values = torch.tensor
【Pytorch学习】torch.mm()&torch.matmul()和torch.mul()以及torch.spmm()
sxl的博客
04-07 1万+
做深度学习过程中免不了使用pytorch，刚上手可能觉得pytorch是一门新语言。过一段时间之后发现其用法跟python极其相似，python的很多函数名都可以用到torch中。当然也有一些不同，毕竟张量的计算可以用GPU啊。是矩阵a和b矩阵相乘，比如a的维度是(1, 2)，b的维度是(2, 3)，返回的就是(1, 3)的矩阵。当然，对于二维的效果等同于torch.mm（）。关于torch.matmul的详细介绍可以看。torch.matmul()的用法比torch.mm更高级，的维度必须相等，比如。
Torch中的矩阵相乘分类
张小白的博客
01-18 9409
矩阵相乘在torch中的几种情况 1、矩阵逐元素(Element-wise)乘法 torch.mul(mat1, other) mat和other可以是标量也可以是任意维度的矩阵，只要满足最终相乘是可以broadcast的即可,即该操作是支持broadcast操作的。只要mat1与other满足broadcast条件，就可可以进行逐元素相乘 。 # 生成指定张量 c = torch.Tensor([[1, 2, 3], [4, 5 ,6]]) print(c.shape) # 2*3 print(c)
【图神经网络】图神经网络框架DGL1.0初体验
小哲的博客
03-16 1024
DGL 1.0 总结了过去三年学术界或工业界对图深度学习和图神经网络（GNN）技术的各类需求。从最先进模型的学术研究到将 GNN 扩展到工业级应用，DGL 1.0 为所有用户提供全面且易用的解决方案，以更好的利用图机器学习的优势。
【学习笔记】图神经网络库 DGL 入门教程（backend pytorch）
热门推荐
一只热爱奔跑的程序猿
12-06 4万+
dgl库笔记 DGL官方文档 目录dgl库笔记1 DGL的安装2 DGL的后端3 一个有趣的入门示例3.1 从"Zachary's karate club" Problem讲起 1 DGL的安装 DGL官方文档 的安装方法似乎有点繁琐, 直接下载wheel文件安装即可; 非CUDA版本的dgl库, 去清华镜像dgl仓库 下载对应版本的whl文件直接用pip install安装即可; CUDA版本的dgl库, 目前有五种不同的dgl库对应不同的CUDA版本: 清华镜像dgl-cu90仓库 ;
return torch._sparse_mm(mat1, mat2)
04-05
As an AI language model, I cannot execute code directly. However, I can explain what the code does. The code `torch._sparse_mm(mat1, mat2)` performs a matrix multiplication of two sparse matrices `mat1` and `mat2` using PyTorch library. Sparse matrices are matrices that have a large number of zero elements. These matrices are used to save memory and computational time when working with large datasets. The PyTorch library provides many functions to work with sparse matrices, including matrix multiplication (`torch._sparse_mm()`), addition, subtraction, and more. The function `torch._sparse_mm(mat1, mat2)` takes two sparse matrices `mat1` and `mat2` and returns their matrix multiplication.

“相关推荐”对你有帮助么？

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值