pytorch---基础篇（常用函数）

说明：大部分关于张量的函数torch.function()都可以使用tensor.function()

1. 张量

1.1 张量基础知识

• pytorch中张量的类型
  
• tensor = tensor.half():将原始张量转换为半精度类型
• tensor = tensor.float():将原始张量转换为float类型
• tensor = tensor.double():将原始张量转换为double类型
• tensor = tensor.short():将原始张量转换为short类型
• tensor = tensor.int():将原始张量转换为int类型
• tensor = tensor.long()将原始张量转换为long类型
• dim=0表示列，第一维；dim=1表示行，第二维

1.2 张量的属性

• tensor.ndim：返回矩阵的纬度数量（有几维）
• tensor.dtype：返回tensor的类型
• tensor.shape:返回torch.Size类型的tensor的size，可以用下标访问。区别于tensor.size()是tensor的方法，返回torch.Size类型，只能通过传参访问tensor纬度。

x = torch.arange(24).reshape(2,3,4)
print(x.shape)
print(x.shape[0])

# torch.Size([2, 3, 4])
# 2

1.3 常用API

• torch.numel(input)->int：返回张量中元素的个数

x = torch.arange(11)
x.numel()  # 11
torch.numel(x)  # 11

• tensor.item()：返回tensor的数值类型（标量）。注意只有张量是一个数的时候才能使用这个方法。
• tensor.size()：返回tensor的size，可以传入dim参数，则返回指定纬度的尺寸

x = torch.arange(24).reshape(2,3,4)
print(x.size())
# torch.Size([2, 3, 4])
print(x.size(1))
# 3

2. 创建操作

• torch.eye(n, m=None, out=None)：返回一个$m\ast n$张量，对角线位置全1，其它位置全0。如果m没指定，则默认为$n\ast n$

>>> torch.eye(3)
 1  0  0
 0  1  0
 0  0  1
[torch.FloatTensor of size 3x3]

• torch.ones(*sizes, out=None) → Tensor：返回全是1的张量

>>> torch.ones(2, 3)
 1  1  1
 1  1  1
[torch.FloatTensor of size 2x3]

• torch.zeros(*sizes, out=None) → Tensor:返回一个全为标量 0 的张量

• torch.rand(*sizes, out=None) → Tensor:返回一个张量，包含了从区间[0,1)的均匀分布中抽取的一组随机数，形状由可变参数sizes 定义

>>> torch.rand(4)
 0.9193
 0.3347
 0.3232
 0.7715
[torch.FloatTensor of size 4]

>>> torch.rand(2, 3)
 0.5010  0.5140  0.0719
 0.1435  0.5636  0.0538
[torch.FloatTensor of size 2x3]

• torch.randn(*sizes, out=None) → Tensor:返回一个张量，包含了从标准正态分布(均值为0，方差为 1，即高斯白噪声)中抽取一组随机数，形状由可变参数sizes定义。
• torch.arange(start, end, step=1, out=None) → Tensor:返回一个1维张量，长度为 floor((end−start)/step)。包含从start到end，以step为步长的一组序列值(默认步长为1)。

>>> torch.arange(1, 4)
 1
 2
 3
[torch.FloatTensor of size 3]

>>> torch.arange(1, 2.5, 0.5)
 1.0000
 1.5000
 2.0000
[torch.FloatTensor of size 3]

• torch.randperm(n, out=None) → LongTensor:给定参数n，返回一个从0 到n -1 的随机整数排列。

• torch.tensor(obj)→ Tensor：将obj转换为tensor类型

input = [
    [2, 3, 4, 5, 0, 0],
    [1, 4, 3, 0, 0, 0],
    [4, 2, 2, 5, 7, 0],
    [1, 0, 0, 0, 0, 0]
]
input = torch.tensor(input)

• torch.linspace(start, end, steps=100, out=None) → Tensor：返回一个1维张量，包含在区间start 和 end 上均匀间隔的steps个点。 输出1维张量的长度为steps。

>>> torch.linspace(3, 10, steps=5)
  3.0000
  4.7500
  6.5000
  8.2500
 10.0000
[torch.FloatTensor of size 5]

3. 索引、切片、连接、换位

3.1 张量索引

• torch.gather(input, dim, index, out=None) → Tensor:沿给定轴dim，将输入索引张量index指定位置的值进行聚合.index–longTensor
  对一个3维张量index，输出可以定义为:
  out[i][j][k] = tensor[index[i][j][k]][j][k] # dim=0
  out[i][j][k] = tensor[i][index[i][j][k]][k] # dim=1
  out[i][j][k] = tensor[i][j][index[i][j][k]] # dim=3
  即对于index中的每个元素，它的索引只有指定纬度变成指定的值，其他纬度的索引位置不变

input = [
    [2, 3, 4, 5, 0, 0],
    [1, 4, 3, 0, 0, 0],
    [4, 2, 2, 5, 7, 0],
    [1, 0, 0, 0, 0, 0]
]
input = torch.tensor(input)
length = torch.LongTensor([[3],[2],[5],[1]])
out = torch.gather(input, 1, length-1)
# tensor([[4],
#         [4],
#         [7],
#         [1]])

t1 = torch.LongTensor([[3,2],[2,2],[4,2],[0,2]])
input.gather(1,t1)
# tensor([[5, 4],
#         [3, 3],
#         [7, 2],
#         [1, 0]])

t = torch.tensor([[1, 2], [3, 4]])
torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
# tensor([[ 1,  1],
#         [ 4,  3]])

3.2 张量切片

• tensor[a:b;c:d]:切片操作，取tensor的a至b行中的c至d列，得到的区间是连续的。tensor[:,i,:]表示取第二维第i个数据，tensor可以是大于等于3维的数据，当取某一维时，前面的不能省，后面可以省。

• torch.narrow(input, dim, start, length) → Tensor:在指定纬度进行一定长度的切片，在dim纬度，从start到start+len-1进行切片

x = torch.randn(3,4)
torch.narrow(x,0,0,2)
# tensor([[ 1.1522, -0.5676,  1.7639,  0.3201],
#         [ 0.4096, -1.7144,  0.8206,  1.6934]])

torch.narrow(x,1,2,2)
# tensor([[ 1.7639,  0.3201],
#         [ 0.8206,  1.6934],
#         [ 0.0301, -0.0711]])

• torch.unbind(tensor, dim=0)[source]:在指定纬度进行切片，返回一个元组，包含了沿着指定维切片后的各个切片

x = torch.randn(2,3,4)
torch.unbind(x,0)
# (tensor([[-0.9200, -0.1279, -1.1582, -0.3785],
#          [ 0.5040, -1.6947, -1.8821, -0.3223],
#          [ 0.4511, -0.4633,  1.4185, -1.6141]]),
#  tensor([[-0.3083, -0.1035, -0.7954, -0.9308],
#          [-0.6792,  0.6268,  0.4888,  1.4539],
#          [-1.3184,  1.4105, -0.0822,  0.6998]]))
torch.unbind(x,1)
# (tensor([[ 0.2602, -0.5469,  1.2912, -0.6746],
#          [ 0.9159, -1.8954, -0.2911,  0.0839]]),
#  tensor([[-1.1844,  1.3560,  0.1255, -1.0270],
#          [-0.5959,  0.0826, -1.5575,  1.4263]]),
#  tensor([[ 0.2494, -0.8061, -1.2643,  1.5658],
#          [-0.4318,  0.4239,  0.9455, -1.4922]]))

• torch.index_select(input, dim, index, out=None) → Tensor:沿着指定维度对输入进行切片，取index中指定的相应项(index为一个LongTensor)，然后返回到一个新的张量， 返回的张量与原始张量_Tensor_有相同的维度(在指定轴上)。可以不连续甚至不是顺序的

>>> x = torch.randn(3, 4)
>>> x
 1.2045  2.4084  0.4001  1.1372
 0.5596  1.5677  0.6219 -0.7954
 1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 3x4]

>>> indices = torch.LongTensor([0, 2])
>>> torch.index_select(x, 0, indices)
 1.2045  2.4084  0.4001  1.1372
 1.3635 -1.2313 -0.5414 -1.8478
[torch.FloatTensor of size 2x4]

>>> torch.index_select(x, 1, indices)
 1.2045  0.4001
 0.5596  0.6219
 1.3635 -0.5414
[torch.FloatTensor of size 3x2]

3.3 张量拼接

• torch.cat(inputs, dimension=0) → Tensor:在给定维度上对输入的张量序列seq 进行连接操作,需要满足除了指定纬度以外的其他纬度长度，且不能广播。返回的张量在指定纬度上长度变为输入张量相应纬度长度之和，其他纬度长度不变。

x = torch.rand(2,3,4)
y = torch.rand(2,3,5)
o1 = torch.cat([x,y],dim=-1)
print(o1.shape)  # torch.Size([2, 3, 9]

# 错误示例
z = torch.rand(2,1,5)
o2 = torch.cat([x,z],dim=-1)
# RuntimeError: Sizes of tensors must match except in dimension 2. Got 3 and 1 in dimension 1

• torch.stack(sequence, dim=0):沿着一个新维度对输入张量序列进行连接（会增加一个新的纬度）。 序列中所有的张量都需要满足为相同形状。

import torch
a = torch.arange(12).reshape(3,4)
b = torch.arange(12,24).reshape(3,4)
c = torch.arange(24,30).reshape(3,2)
torch.stack([a,b])
# tensor([[[ 0,  1,  2,  3],
#          [ 4,  5,  6,  7],
#          [ 8,  9, 10, 11]],

#         [[12, 13, 14, 15],
#          [16, 17, 18, 19],
#          [20, 21, 22, 23]]])

torch.cat([a,c],dim=1)
# tensor([[ 0,  1,  2,  3, 24, 25],
#         [ 4,  5,  6,  7, 26, 27],
#         [ 8,  9, 10, 11, 28, 29]])

3.4 张量分块

• torch.chunk(tensor, chunks, dim=0):在给定维度(轴)上将输入张量进行分块儿。chunks (int) – 分块的个数

x = torch.arange(11)
y = torch.chunk(x,3)
y = x.chunk(3)
# (tensor([0, 1, 2, 3]), tensor([4, 5, 6, 7]), tensor([ 8,  9, 10]))

• torch.split(tensor, split_size, dim=0):将输入张量分割成指定shape的块，区别于chunk指定的是分块的个数，split指定的是分块的形状

>>> a = torch.arange(10).reshape(5,2)
>>> a
tensor([[0, 1],
        [2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])
        
>>> torch.split(a, 2)
(tensor([[0, 1],
         [2, 3]]),
 tensor([[4, 5],
         [6, 7]]),
 tensor([[8, 9]]))
 
>>> torch.split(a, [1,4])
(tensor([[0, 1]]),
 tensor([[2, 3],
         [4, 5],
         [6, 7],
         [8, 9]]))

3.5压缩与扩维

• torch.squeeze():降维
  （1）squeeze(a)就是将a中所有为1的维度删掉。不为1的维度没有影响。
  （2）a.squeeze(N) 就是去掉a中指定的维数为一的维度。
• torch.unsqueeze()：对纬度扩充
  b=torch.unsqueeze(a，N) a就是在a中指定位置N加上一个维数为1的维度
  参考链接

3.6 纬度扩充

• Tensor.expand(*sizes) → Tensor:返回原有tensor扩展到指定的纬度长度的新视图。这里的参数表示最终扩展的纬度长度。-1表示该纬度的长度不变。
  ps：返回的是新视图，在内存中是不连续的。且如果纬度个数增加了只能在前面增加；此外，expand函数只能对纬度长度为1的纬度进行扩充

x = torch.arange(4)
y = x.expand(2,3,-1)
print(x.shape,y.shape)
# torch.Size([4]) torch.Size([2, 3, 4])

print(x,'\n',y)
# tensor([0, 1, 2, 3]) 
# tensor([[[0, 1, 2, 3],
#      [0, 1, 2, 3],
#      [0, 1, 2, 3]],

#     [[0, 1, 2, 3],
#      [0, 1, 2, 3],
#      [0, 1, 2, 3]]])

• Tensor.repeat(*sizes) → Tensor：根据参数在该纬度重复指定的次数，且参数的个数不能小于tensor的纬度。这里的参数表示延该纬度复制的次数。例如参数为2,3,4，则表示在第1，2，3维分别重复2，3，4次.
  ps：repeat函数拷贝原有的数据，不同于expand只是返回一个新视图。如果参数的数量大于原始张量的纬度数量，则先将原始张量执行unsqueeze操作（且在第一维扩充）到与参数纬度的数目。在下面这个例子，x：（4,）-> （1,1,4）->（12, 13, 4*4）。
  此外，repeat只是tensor的方法，不是torch的方法，即不能使用torch.repeat()，只有tensor.repeat()

x = torch.arange(4)
y = x.repeat(2,3,4)
print(x.shape,y.shape)
# torch.Size([4]) torch.Size([2, 3, 16])

print(x,'\n',y)
# tensor([0, 1, 2, 3]) 
# tensor([[[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3],
#      [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3],
#      [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]],

#     [[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3],
#      [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3],
#      [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]])

3.7 纬度置换

• torch.transpose(input, dim0, dim1, out=None) → Tensor:交换input的维度dim0和dim1。

x = torch.randn(4,3,2)
x.transpose(0,2)
# tensor([[[ 1.4651,  0.4556, -2.4765,  0.0360],
#          [ 0.3542,  1.3971,  0.7538, -0.0903],
#          [ 0.7990, -1.3353,  0.1951, -0.8441]],

#         [[ 0.5109,  0.5943, -0.0418,  2.0164],
#          [-2.2539, -1.3724, -0.6792, -0.6619],
#          [-0.9598,  2.1428,  0.3904, -1.2091]]])

• torch.t(input, out=None) → Tensor:对二维张量进行转置。理论上要求input的纬度数<=2，当张量为一维时，返回原始张量。 可以被视为函数transpose(input, 0, 1)的简写函数。

>>> x = torch.randn(2, 3)
>>>> torch.t(x)
 0.4834 -0.1300
 0.6907  0.5295
 1.3417  0.2321
[torch.FloatTensor of size 3x2]

• torch.permute(input, dims) → Tensor:对矩阵纬度进行排列

>>> x = torch.randn(2, 3, 5)
>>> x.size()
torch.Size([2, 3, 5])
>>> torch.permute(x, (2, 0, 1)).size()
torch.Size([5, 2, 3])

3.8 纬度变形

• torch.reshape(input, shape) → Tensor:改变矩阵的size，-1表示该纬度的长度由其他纬度决定

>>> a = torch.arange(4.)
>>> torch.reshape(a, (2, 2))
tensor([[ 0.,  1.],
        [ 2.,  3.]])
>>> b = torch.tensor([[0, 1], [2, 3]])
>>> torch.reshape(b, (-1,))
tensor([ 0,  1,  2,  3])

• Tensor.view(*shape) → Tensor:作用于用法同reshape。区别如下：
  view:返回具有新形状的张量。返回的张量将与原始张量共享基础数据。且只能在连续的张量上运行。如果对 tensor 调用过 transpose, permute 等操作的话会使该 tensor 在内存中变得不再连续，此时就不能再调用 view 函数
  reshape: torch.reshape可能返回原始张量的副本或视图->不确定的。可以在连续和非连续的张量上运行。

• torch.flatten(input, start_dim=0, end_dim=- 1) → Tensor：将tensor从start_dim到end_dim中的数据进行展平。

tensor t : 3*2*2
torch.flatten(t) # (12,)  默认形式：展平成一维
torch.flatten(t，1)  # (3, 4) 从第1维开始展平
torch.flatten(t， 0， 1)    #(6,2) 从第0维到第1维展平

4. 随机抽样

• torch.manual_seed(seed) → torch._C.Generator:设定生成随机数的种子.设置CPU生成随机数的种子，方便下次复现实验结果。seed–int

import torch
torch.manual_seed(0)
print(torch.rand(1)) 
# 每次生成的数字都是一样的

• torch.bernoulli(input, out=None) → Tensor：从伯努利分布中抽取二元随机数(0 或者 1)。输入是抽样等于1的概率，返回是抽样结果（0或1）

>>> a = torch.Tensor(3, 3).uniform_(0, 1) # generate a uniform random matrix with range [0, 1]
>>> a
 0.7544  0.8140  0.9842
 0.5282  0.0595  0.6445
 0.1925  0.9553  0.9732
[torch.FloatTensor of size 3x3]

>>> torch.bernoulli(a)
 1  1  1
 0  0  1
 0  1  1
[torch.FloatTensor of size 3x3]

>>> a = torch.ones(3, 3) # probability of drawing "1" is 1
>>> torch.bernoulli(a)
 1  1  1
 1  1  1
 1  1  1
[torch.FloatTensor of size 3x3]

>>> a = torch.zeros(3, 3) # probability of drawing "1" is 0
>>> torch.bernoulli(a)
 0  0  0
 0  0  0
 0  0  0
[torch.FloatTensor of size 3x3]

• torch.multinomial(input, num_samples,replacement=False, out=None) → LongTensor:

5. 序列化（模型的保存和加载）

• torch.save(obj, f, pickle_module=, pickle_protocol=2):保存一个对象到一个硬盘文件上。一般用于保存模型或模型参数
  obj – 保存对象
  f － 类文件对象 (返回文件描述符)或一个保存文件名的字符串
  pickle_module – 用于pickling元数据和对象的模块
  pickle_protocol – 指定pickle protocal 可以覆盖默认参数

#方式一：保存整个模型，消耗存储较大
torch.save(model1,'model1.pt')

#方式二：保存模型参数，消耗存储较小
torch.save(model2.state_dict(),'model2.pt')

• torch.load(f, map_location=None, pickle_module=):从磁盘文件中读取一个通过torch.save()保存的对象
  f – 类文件对象 (返回文件描述符)或一个保存文件名的字符串
  map_location – 一个函数或字典规定如何remap存储位置
  pickle_module – 用于unpickling元数据和对象的模块 (必须匹配序列化文件时的pickle_module )

# 方式一加载整个模型
model1 = torch.load('model1.pt')

# 方式二加载模型参数
model2 = Model()
model2.load_state_dict(torch.load('model2.pt'))

torch.load('tensors.pth')
# 把所有的张量加载到CPU中
torch.load('tensors.pth', map_location=lambda storage, loc: storage)
# 把所有的张量加载到GPU 1中
torch.load('tensors.pth', map_location=lambda storage, loc: storage.cuda(1))
# 把张量从GPU 1 移动到 GPU 0
torch.load('tensors.pth', map_location={'cuda:1':'cuda:0'})

6. 并行化

6.1 介绍

（1）每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。

在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。

在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。之后，各进程用该梯度来更新参数。由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。

而在 DataParallel 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。

（2）每个进程包含独立的解释器和 GIL
由于每个进程拥有独立的解释器和 GIL，消除了来自单个 Python 进程中的多个执行线程，模型副本或 GPU 的额外解释器开销和 GIL-thrashing ，因此可以减少解释器和 GIL 使用冲突。这对于严重依赖 Python runtime 的 models 而言，比如说包含 RNN 层或大量小组件的 models 而言，这尤为重要。
参考链接

6.2 流程

（1）指定程序可见的GPU

• 方法一：在程序开头指定

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3,4"   #指定程序可见的GPU

注意，这两句一定要在程序入口的最开始指定（一定要在import torch之前），否则会失效。
此外，这句话生效之后，此后所有与GPU编号有关的数字都是相对编号。例如这里的GPU1编号为0，以此类推，且主卡为编号为0的卡。

• 方法二：在终端指定

CUDA_VISIBLE_DEVICES=0,1,2,3 python train_debug.py

这里的XXXXX需要换成自己的端口号，nproc_per_node表示表示每个节点上有多少个进程，一般填写参与训练的GPU的数量。

• 多GPU训练流程

（2）引入local_rank参数。
由torch.distributed.launch提供，用于指定每个GPU在本地的rank。

parser.add_argument('--local_rank', type=int, default=-1)

（3）初始化进程组和包

dist.init_process_group(backend='nccl')

一般建议如果使用CPU，参数设置‘gloo’，这里为GPU所以建议设置为’nccl’，性能更好。

（4）配置每个进程的GPU
我们需要为每个GPU启动一个进程。每个进程需要知道自己运行在哪个GPU上，以及自身在所有进程中的序号。

torch.cuda.set_device(args.local_rank)  
device = torch.device(f'cuda:{args.local_rank}')

（5）使用DistributedSampler
为各个进程切分数据

dataset = xxx
data_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
train_loader = torch.utils.data.DataLoader(dataset=dataset,
                         batch_size=batch_size,
                         sampler=data_sampler)

（6）封装之前要把模型、数据移到对应的gpu

model.to(device)

（7）将模型封装
在背后它会支持梯度的All-Reduce操作

model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[local_rank])

这里的divice_ids表示参与训练的GPU编号（相对编号）

（8）在每个 epoch 开始时调用 set_epoch() 方法
在分布式模式下，需要在每个 epoch 开始时调用 set_epoch() 方法，然后再创建 DataLoader 迭代器，以使 shuffle 操作能够在多个 epoch 中正常工作。 否则，dataloader迭代器产生的数据将始终使用相同的顺序。

sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None),
                    sampler=sampler)
for epoch in range(start_epoch, n_epochs):
    if is_distributed:
        sampler.set_epoch(epoch)
    train(loader)

不过为什么我看到一份代码是这样写的，并不是每一个epoch都会调用set_epoch()，这又是为啥呀？待补充…

if self.multi_gpu and self.local_rank<=0:
    self.train_loader.sampler.set_epoch(epoch)

（9）保存模型
保存模型需要加上module。由于每个进程的参数是一致的，所以只需要在一个进程中保存模型参数即可

model.module.state_dict()

（10）torch.distributed.launch （可在终端或参数中指定）
torch.distributed.launch 会给模型分配一个args.local_rank的参数

python -m torch.distributed.launch train_debug.py

python -m torch.distributed.launch main.py相当于使用torch.distributed.launch.py来运行我们的main.py，其中torch.distributed.launch会向我们的运行程序传递一些变量，包括nproc_per_node参数。该参数将会引导当前主机创建nproc_per_node个进程，每个进程会独立执行训练脚本。同时，每个进程会被分配一个local_rank参数来表示进程在当前主机（主机的参数是rank，如果是一个主机，就默认为0）上的编号，用以合理分配和调度本地的GPU资源（这也是为什么需要torch.cuda.set_device(args.local_rank)设定默认的GPU，因为每个进程需要在一个独立的GPU上）。在实际应用中，DDP会自动帮助我们将模型从local_rank=0扩展到其他进程状态。
参考链接

6.3 评估与测试

由于每个进程只是计算了一个batch中的部分数据，对于模型的参数，DDP会自动对其进行all_reduce操作。如果想要统计其他数据，则需要手动进行all_reduce。

• torch.distributed.all_reduce(tensor, op=, group=None, async_op=False):不同进程的GPU中关于tensor的数值变为执行前所有GPU的和。

>>> # All tensors below are of torch.int64 type.
>>> # We have 2 process groups, 2 ranks.
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_reduce(tensor, op=ReduceOp.SUM)
>>> tensor
tensor([4, 6]) # Rank 0
tensor([4, 6]) # Rank 1

• torch.distributed.reduce(tensor, dst, op=, group=None, async_op=False):与all_reduce类似，只不过最终结果只在rank=dst的GPU上复制

• torch.distributed.all_gather(tensor_list, tensor, group=None, async_op=False:从所有设备收集指定的input_tensor并将其放置在所有设备上的tensor_list变量中。执行后tensor_list将成为所有GPU上tensor构成的list。

>>> # All tensors below are of torch.int64 dtype.
>>> # We have 2 process groups, 2 ranks.
>>> tensor_list = [torch.zeros(2, dtype=torch.int64) for _ in range(2)]
>>> tensor_list
[tensor([0, 0]), tensor([0, 0])] # Rank 0 and 1
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> tensor
tensor([1, 2]) # Rank 0
tensor([3, 4]) # Rank 1
>>> dist.all_gather(tensor_list, tensor)
>>> tensor_list
[tensor([1, 2]), tensor([3, 4])] # Rank 0
[tensor([1, 2]), tensor([3, 4])] # Rank 1

• torch.distributed.all_gather_object(object_list, obj, group=None):与all_gather类似，只不过对象由tensor变成python object，注意object需要是picklable的以保证能够被打包。

object_list (list[Any]) – Output list. It should be correctly sized as the size of the group for this collective and will contain the output.
object (Any) – Pickable Python object to be broadcast from current process.
group (ProcessGroup, optional) – The process group to work on. If None, the default process group will be used. Default is None.

• torch.distributed.gather(tensor, gather_list=None, dst=0, group=None, async_op=False):从所有设备收集指定的tensor并将它们放置在gather_list中的dst设备上。
  

7. 数学操作

7.1 elementwise

7.1.1 数学运算

• +-*/**：张量的对应元素操作，或者张量与数值操作，可以直接使用相应符号。但是要保证两个张量的size是一致的

x = torch.rand(2,3)
y = torch.randn(2,3)
-x
x+y  x+z
x-y  x-z
x*y  x*z
x/y  x/z
x%y  x%z
x**y  x**z
...

相应的torch函数如下（支持广播）：

• torch.neg(input, out=None) → Tensor:逐元素取负
• torch.add(input, value, out=None) , `torch.add(input, value=1, other, out=None):相加，相减第二个取负即可
• torch.mul(input, value, out=None) , torch.mul(input, other, out=None)：相乘
• torch.div(input, value, out=None) , torch.div(input, other, out=None)：相除
• torch.fmod(input, divisor, out=None) → Tensor:计算除法余数。 除数与被除数可能同时含有整数和浮点数。此时，余数的正负与被除数相同。
• torch.remainder(input, divisor, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的除法余数。 除数与被除数可能同时包含整数或浮点数。余数与除数有相同的符号。
  参数： - input (Tensor) – 被除数 - divisor (Tensor or float) – 除数，一个数或与被除数相同类型的张量 - out (Tensor, optional) – 输出张量
• torch.pow(input, exponent, out=None):对输入input的按元素求exponent次幂值，并返回结果张量。 幂值exponent 可以为单一 float 数或者与input相同元素数的张量。
• torch.exp(tensor, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的指数。
• torch.log(input, out=None) → Tensor:计算input 的自然对数
• torch.log1p(input, out=None) → Tensor:计算 input+1的自然对数 yi=log(xi+1)
• torch.reciprocal(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的倒数，即 1.0/x。
• torch.sqrt(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的平方根。
• torch.rsqrt(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的平方根倒数
• torch.abs(input, out=None) → Tensor:计算输入张量的每个元素绝对值

>>> torch.abs(torch.FloatTensor([-1, -2, 3]))
FloatTensor([1, 2, 3])

7.1.2 近似

• torch.ceil(input, out=None) → Tensor：天花板函数，输入input张量每个元素向上取整, 即取不小于每个元素的最小整数，并返回结果到输出。

• torch.floor(input, out=None) → Tensor:地板函数， 返回一个新张量，包含输入input张量每个元素的floor，即不小于元素的最大整数。

• torch.round(input, out=None) → Tensor:返回一个新张量，将输入input张量每个元素舍入到最近的整数。

7.1.3 截断

• torch.trunc(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的截断值(标量x的截断值是最接近0的整数，其比x更接近零。简而言之，有符号数的小数部分被舍弃)
• torch.frac(tensor, out=None) → Tensor:返回每个元素的分数（小数）部分（带符号）。

>>> torch.frac(torch.Tensor([1, 2.5, -3.2])
torch.FloatTensor([0, 0.5, -0.2])

• torch.clamp(input, min, max, out=None) → Tensor:将输入input张量每个元素的夹紧到区间 [min,max]，并返回结果到一个新张量。
           | min, if x_i < min
  y_i = | x_i, if min <= x_i <= max
           | max, if x_i > max

>>> a = torch.randn(4)
>>> a
 1.3869
 0.3912
-0.8634
-0.5468

>>> torch.clamp(a, min=-0.5, max=0.5)
 0.5000
 0.3912
-0.5000
-0.5000

• torch.clamp(input, *, min, out=None) → Tensor:将输入input张量每个元素的限制到不小于min ，并返回结果到一个新张量

torch.clamp(a, min=0.5)

• torch.clamp(input, *, max, out=None) → Tensor:将输入input张量每个元素的限制到不大于max ，并返回结果到一个新张量。
• torch.lerp(start, end, weight, out=None):对两个张量以start，end做线性插值， 将结果返回到输出张量。weight是float类型$ou{t}_i=star{t}_i+weight\ast (en{d}_i-star{t}_i)$

>>> start = torch.arange(1, 5)
>>> end = torch.Tensor(4).fill_(10)
>>> start
 1
 2
 3
 4
[torch.FloatTensor of size 4]

>>> end
 10
 10
 10
 10
[torch.FloatTensor of size 4]

>>> torch.lerp(start, end, 0.5)
 5.5000
 6.0000
 6.5000
 7.0000
[torch.FloatTensor of size 4]

7.1.4 激活函数

• torch.sigmoid(input, out=None) → Tensor：返回一个新张量，包含输入input张量每个元素的sigmoid值。
• torch.sign(input, out=None) → Tensor:符号函数：返回一个新张量，包含输入input张量每个元素的正负。

>>> a = torch.randn(4)
>>> a
-0.6366
 0.2718
 0.4469
 1.3122
[torch.FloatTensor of size 4]

>>> torch.sign(a)
-1
 1
 1
 1
[torch.FloatTensor of size 4]

7.1.5 三角函数

• torch.sin(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的正弦。
• torch.cos(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的余弦。
• torch.tan(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的正切。
• torch.asin(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的反正弦函数
• torch.acos(input, out=None) → Tensor:返回一个新张量，包含输入张量每个元素的反余弦
• torch.atan(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的反正切函数
• torch.atan2(input1, input2, out=None) → Tensor:返回一个新张量，包含两个输入张量input1和input2的反正切函数???
• torch.cosh(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的双曲余弦。
• torch.sinh(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的双曲正弦。
• torch.tanh(input, out=None) → Tensor:返回一个新张量，包含输入input张量每个元素的双曲正切。

7.2 矩阵操作

7.2.1 矩阵乘法

• torch.dot(tensor1, tensor2) → float：向量点乘。只允许输入是两个向量（一维），返回两个向量点乘的结果
• torch.mv(input, vec, *, out=None) → Tensor:矩阵和向量相乘。input为$n\ast m$，vec为$(m,)$
• torch.mm(input, mat2, *, out=None) → Tensor:矩阵相乘，input为$n\ast m$，mat2为$m\ast p$，结果为$n\ast p$
• torch.matmul(input, other, *, out=None) → Tensor :支持多种纬度的张量相乘。如果第一个参数为1维，第二个参数纬度大于2，则第一个参数会在其维数之前加1（d->1*d）以进行批处理矩阵乘法；反之类似

>>> # vector x vector
>>> tensor1 = torch.randn(3)
>>> tensor2 = torch.randn(3)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([])
>>> # matrix x vector
>>> tensor1 = torch.randn(3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([3])
>>> # batched matrix x broadcasted vector
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3])
>>> # batched matrix x batched matrix
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(10, 4, 5)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3, 5])
>>> # batched matrix x broadcasted matrix
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(4, 5)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3, 5])

详细介绍戳torch中的乘法函数

• torch.bmm(input, mat2, *, out=None) → Tensor:batched matrix相乘。input和mat2都必须是三维张量，且第一维是batch size，剩下两维需要满足矩阵相乘的要求。例如input shape为$N\ast n\ast m$，mat2 shape为$N\ast m\ast p$，则output shape为$N\ast n\ast p$

>>> input = torch.randn(10, 3, 4)
>>> mat2 = torch.randn(10, 4, 5)
>>> res = torch.bmm(input, mat2)
>>> res.size()
torch.Size([10, 3, 5])

7.2.2 累计操作

• torch.sum(input) → float:返回输入张量input 所有元素的和。
• torch.sum(input, dim, out=None) → Tensor:返回输入张量给定维度上每行的和。 输出形状与输入相同，除了给定维度上为1.
• torch.prod(input) → float:返回输入张量input 所有元素的积
• torch.prod(input, dim, out=None) → Tensor:返回输入张量给定维度上每行的积。 输出形状与输入相同，除了给定维度上为1.
• torch.cumprod(input, dim, out=None) → Tensor:返回输入沿指定维度的累积积。例如，如果输入是一个N 元向量，则结果也是一个N 元向量，第i 个输出元素值为yi=x1∗x2∗x3∗…∗xi

x = torch.arange(2,8).reshape(2,3)
# tensor([[2, 3, 4],
#         [5, 6, 7]])
torch.cumprod(x,0)
# tensor([[ 2,  3,  4],
#         [10, 18, 28]]

• torch.cumsum(input, dim, out=None) → Tensor:返回输入沿指定维度的累积和。例如，如果输入是一个N元向量，则结果也是一个N元向量，第i 个输出元素值为 yi=x1+x2+x3+…+xi

7.2.3 矩阵统计量

• torch.dist(input, other, p=2, out=None) → Tensor:返回 (input - other) 的 p范数 。p表示所计算的范数，类型为float
  例如二范数的计算公式为${||x||}_2=\sum \sqrt{x_i^2}$
• torch.norm(input, p=2) → float:返回输入张量input 的p 范数。
• torch.mean(input) → float:返回输入张量所有元素的均值。
• torch.mean(input, dim, keepdim=False, *, dtype=None, out=None) → Tensor:返回输入张量给定维度dim上每行的均值。
  dim：对指定的纬度求均值。可以是int或tuple。例如若是两个元素的tuple，则是对二维的数据求均值。返回的结果在指定的纬度上被消解
  keepdim：若为true，则表示返回的结果与原始张量在纬度的个数上保持一致，在dim纬度的长度为1

x = torch.arange(24).reshape(2,3,4).float()
print(x,'\n')
# tensor([[[ 0.,  1.,  2.,  3.],
#          [ 4.,  5.,  6.,  7.],
#          [ 8.,  9., 10., 11.]],

#         [[12., 13., 14., 15.],
#          [16., 17., 18., 19.],
#          [20., 21., 22., 23.]]])

y = torch.mean(x,1)
print(y.shape,'\n',y,'\n')
torch.Size([2, 4]) 
# tensor([[ 4.,  5.,  6.,  7.],
#         [16., 17., 18., 19.]]) 

y = torch.mean(x,1,keepdim=True)
print(y.shape,'\n',y,'\n')
torch.Size([2, 1, 4]) 
# tensor([[[ 4.,  5.,  6.,  7.]],

#         [[16., 17., 18., 19.]]])

y = torch.mean(x,(0,2))
print(y.shape,'\n',y)
# torch.Size([3]) 
# tensor([ 7.5000, 11.5000, 15.5000])

• torch.var(input) → float:返回输入张量所有元素的方差
• torch.var(input, dim, out=None) → Tenso:返回输入张量给定维度上的方差。 输出形状与输入相同，除了给定维度上为1.
• torch.std(input) → float:返回输入张量input 所有元素的标准差。
• torch.std(input, dim, out=None) → Tensor:返回输入张量给定维度上每行的标准差。 输出形状与输入相同，除了给定维度上为1.
• torch.median(input, dim=-1, values=None, indices=None) -> (Tensor, LongTensor):返回一个命名的元组包含$(values,indies)$，其中values为输入张量给定维度每行的中位数，indies是包含中位数的索引的LongTensor。dim值默认为输入张量的最后一维。 输出形状与输入相同，除了给定维度上为1.

x = torch.arange(2,8).reshape(2,3).float()

print(torch.median(x,1))
# torch.return_types.median(
# values=tensor([3., 6.]),
# indices=tensor([1, 1])) 

print(torch.median(x,1).values)
# tensor([3., 6.])

print(torch.median(x,1).indices)
# tensor([1, 1])

• torch.mode(input, dim=-1, values=None, indices=None) -> (Tensor, LongTensor):返回给定维dim上，每行的众数值。 同时返回一个LongTensor，包含众数职的索引。dim值默认为输入张量的最后一维。输出形状与输入相同，除了给定维度上为1.用法与median类似

7.3 other

• torch.tril(input, diagonal=0, *, out=None) → Tensor:返回矩阵的下三角部分，diagonal参数控制对角线的移动。diagonal=0返回一个标准的下三角矩阵，>0表示对角线上移，反之下移

>>> a = torch.randn(3, 3)
>>> a
tensor([[-1.0813, -0.8619,  0.7105],
        [ 0.0935,  0.1380,  2.2112],
        [-0.3409, -0.9828,  0.0289]])
>>> torch.tril(a)
tensor([[-1.0813,  0.0000,  0.0000],
        [ 0.0935,  0.1380,  0.0000],
        [-0.3409, -0.9828,  0.0289]])

>>> b = torch.randn(4, 6)
>>> b
tensor([[ 1.2219,  0.5653, -0.2521, -0.2345,  1.2544,  0.3461],
        [ 0.4785, -0.4477,  0.6049,  0.6368,  0.8775,  0.7145],
        [ 1.1502,  3.2716, -1.1243, -0.5413,  0.3615,  0.6864],
        [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024,  0.0978]])
>>> torch.tril(b, diagonal=1)
tensor([[ 1.2219,  0.5653,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.4785, -0.4477,  0.6049,  0.0000,  0.0000,  0.0000],
        [ 1.1502,  3.2716, -1.1243, -0.5413,  0.0000,  0.0000],
        [-0.0614, -0.7344, -1.3164, -0.7648, -1.4024,  0.0000]])
>>> torch.tril(b, diagonal=-1)
tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.4785,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 1.1502,  3.2716,  0.0000,  0.0000,  0.0000,  0.0000],
        [-0.0614, -0.7344, -1.3164,  0.0000,  0.0000,  0.0000]])

• Tensor.masked_fill_(mask, value):将tensor中mask为true的部分用value填充。支持广播，但是要保证张量的纬度数目一致。即tensor和mask的len(shape)一致，mask每个纬度的长度要么与tensor相同，要么为1

x = torch.rand(2,3,4)
y = torch.rand(2,1,4)
z = y.expand(-1,3,-1)
a = x.masked_fill(y>0.5,10)
b = x.masked_fill(z>0.5,10)
print(a.equal(b))   # True

8. 比较操作

• torch.eq(input, other, out=None) → Tensor:比较元素相等性。第二个参数可为一个数或与第一个参数同类型形状的张量。

>>> torch.eq(torch.Tensor([[1, 2], [3, 4]]), torch.Tensor([[1, 1], [4, 4]]))
1  0
0  1
[torch.ByteTensor of size 2x2]

• torch.equal(tensor1, tensor2) → bool:如果两个张量有相同的形状和元素值，则返回True ，否则 False。

>>> torch.equal(torch.Tensor([1, 2]), torch.Tensor([1, 2]))
True

• torch.ge(input, other, out=None) → Tensor:逐元素比较input和other，即是否 input>=other。other (Tensor or float) – 对比的张量或float值

>>> torch.ge(torch.Tensor([[1, 2], [3, 4]]), torch.Tensor([[1, 1], [4, 4]]))
 1  1
 0  1
[torch.ByteTensor of size 2x2]

• torch.gt(input, other, out=None) → Tensor:逐元素比较input和other ， 即是否input>other .out (Tensor, optional) – 输出张量。必须为ByteTensor或者与第一个参数tensor相同类型。
• torch.le(input, other, out=None) → Tensor：input是否小于等于output
• torch.lt(input, other, out=None) → Tensor:严格小于
• torch.kthvalue(input, k, dim=None, out=None) -> (Tensor, LongTensor):取输入张量input指定维上第k 个最小值。如果不指定dim，则默认为input的最后一维。
  返回一个元组 (values,indices)，其中indices是原始输入张量input中沿dim维的第 k 个最小值下标。
• torch.max(input) → float:返回输入张量所有元素的最大值。
• torch.max(input, dim, max=None, max_indices=None) -> (Tensor, LongTensor):返回输入张量给定维度上每行的最大值，并同时返回每个最大值的位置索引。(values,indies)
• torch.max(input, other, out=None) → Tensor:outi=max(inputi,otheri)
• torch.argmax(input,dim=None,keepdim=None) → LongTensor：返回最大值的索引，即max()返回值的第二项。如果dim为None（默认情况下），则返回tensor展平后的索引。
• torch.min(input) → float:返回输入张量所有元素的最小值。
• torch.min(input, dim, min=None, min_indices=None) -> (Tensor, LongTensor)：用法类似
• torch.min(input, other, out=None) → Tensor:用法类似
• torch.ne(input, other, out=None) → Tensor:逐元素比较input和other ， 即是否 input!=other
• torch.sort(input, dim=None, descending=False, out=None) -> (Tensor, LongTensor):对输入张量input沿着指定维按升序排序。如果不给定dim，则默认为输入的最后一维。如果指定参数descending为True，则按降序排序。
  返回元组 (sorted_tensor, sorted_indices) ， sorted_indices 为原始输入中的下标。

>>> x = torch.randn(3, 4)
>>> sorted, indices = torch.sort(x)
>>> sorted
-1.6747  0.0610  0.1190  1.4137
-1.4782  0.7159  1.0341  1.3678
-0.3324 -0.0782  0.3518  0.4763
[torch.FloatTensor of size 3x4]

>>> indices
 0  1  3  2
 2  1  0  3
 3  1  0  2
[torch.LongTensor of size 3x4]

>>> sorted, indices = torch.sort(x, 0)
>>> sorted
-1.6747 -0.0782 -1.4782 -0.3324
 0.3518  0.0610  0.4763  0.1190
 1.0341  0.7159  1.4137  1.3678
[torch.FloatTensor of size 3x4]

>>> indices
 0  2  1  2
 2  0  2  0
 1  1  0  1
[torch.LongTensor of size 3x4]

• torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor):沿给定dim维度返回输入张量input中 k 个最大值。 如果不指定dim，则默认为input的最后一维。 如果为largest为 False ，则返回最小的 k 个值。
  返回一个元组 (values,indices)，其中indices是原始输入张量input中测元素下标。 如果设定布尔值sorted 为_True_，将会确保返回的 k 个值被排序。

>>> x = torch.arange(1, 6)
>>> torch.topk(x, 3)
(5
 4
 3
[torch.FloatTensor of size 3],
 4
 3
 2
[torch.LongTensor of size 3])

>>> torch.topk(x, 3, 0, largest=False)
(1
 2
 3
[torch.FloatTensor of size 3],
 0
 1
 2
[torch.LongTensor of size 3])

9. 稀疏矩阵（张量）

• 简要介绍：如果在一个高维张量中，只有少部分位置是有值的，其他位置的值为空或为0或不重要。那么为了减少存储以及加快计算可以使用稀疏存储方式（例如COO的存储是一个三元组，分别包括下标，值以及张量尺寸）

9.1 创建

• torch.sparse.FloatTensor(i, v, torch.Size()):创建一个COO类型的稀疏矩阵。需要三个参数，分别是张量中数值的下标（按维度排列的list），数值，以及张量的size。

i = torch.LongTensor([[0, 1, 1],   #第一维
                      [2, 0, 2]])    #第二维
v = torch.FloatTensor([3, 4, 5])
sp = torch.sparse.FloatTensor(i, v, torch.Size([2,3]))
print(adj)

# tensor(indices=tensor([[0, 1, 1],
#                        [2, 0, 2]]),
#        values=tensor([3., 4., 5.]),
#        size=(2, 3), nnz=3, layout=torch.sparse_coo)

• 高维稀疏张量

num_edge_type = 2
n_node = 4
index = torch.tensor([[0, 1],  # 第一维
                     [0, 2],   # 第二维
                     [1, 3]])  #第三维
value = torch.tensor([1.0, 1.0])
shape = torch.Size([num_edge_type, n_node, n_node])

# 创建稀疏矩阵
sp = torch.sparse.FloatTensor(index, value, shape)
print('sp:\n', sp)
# 转化为稠密矩阵
sp_dense = sp.to_dense()
print('sp_dense:\n', sp_dense)

# 输出
# sp:
#  tensor(indices=tensor([[0, 1],
#                        [0, 2],
#                        [1, 3]]),
#        values=tensor([1., 1.]),
#        size=(2, 4, 4), nnz=2, layout=torch.sparse_coo)
# sp_dense:
#  tensor([[[0., 1., 0., 0.],
#          [0., 0., 0., 0.],
#          [0., 0., 0., 0.],
#          [0., 0., 0., 0.]],
#         [[0., 0., 0., 0.],
#          [0., 0., 0., 0.],
#          [0., 0., 0., 1.],
#          [0., 0., 0., 0.]]])

9.2 相互转化

• to_dense():将稀疏矩阵转化为稠密矩阵。

sp.to_dense()
# tensor([[0., 0., 3.],
#         [4., 0., 5.]])

9.2 操作

• torch.sparse.mm(mat1: torch.Tensor, mat2: torch.Tensor) → torch.Tensor：仅支持二维矩阵相乘。mat1必须是sparse matrix，mat2必须是dense matrix。

m2 = torch.arange(1,10).reshape(3,3).float()
torch.sparse.mm(sp,m2)

# tensor([[21., 24., 27.],
#         [39., 48., 57.]])

• torch.spmm(mat1,mat2):与torch.sparse.mm()用法相同，应该就是torch.sparse.mm()。
• sp.sparse_resize_()：改变稀疏张量的尺寸

new_nodes_size = 6    # 新的图的节点个数 
sp.sparse_resize_(size=(2, new_nodes_size, new_nodes_size), sparse_dim=3, dense_dim=0)

• torch.sparse.sum（input：torch.Tensor，dim：Optional[Tuple[int]] = None，dtype：Optional[int] = None ） → torch.Tensor:返回值为指定维度的和，通过dim参数来指定维度，dim为一个Tuple

x = torch.sparse.sum(a, dim=[0])  # 对列求和
tensor([9., 6., 3.])
x = torch.sparse.sum(a, dim=[1])  # 对行求和
tensor([ 3., 15.])

更多操作见官方文档
参考链接：
pytorch稀疏矩阵
图神经网络中的稀疏矩阵