Pytorch常用函数集锦

Pytorch的函数使用频率也符合 20-80 法则：

• 20%的函数高频使用
• 80%的函数低频使用
  本文总结一下高频使用的那20%函数

第一，查看版本信息

# 查看版本信息
print(torch.__version__)    # torch version
print(torch.version.cuda)   # cuda version
print(torch.backends.cudnn.version()) # cudnn version
print(torch.cuda.get_device_name(0))  # GPU name

1.5.0
10.1
7604
GeForce GTX 1080 Ti

第二，固定随机种子, 以便复现训练结果;指定程序运行在特定GPU卡上;判断是否有CUDA支持

# 固定随机种子, 以便复现训练结果
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

# 指定程序运行在特定GPU卡上
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1' #指定在第0,1块GPU上运行

# 判断是否有CUDA支持
print(torch.cuda.is_available())

第三，设置为cuDNN benchmark模式,Benchmark模式会提升计算速度

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
# 训练被中断，需要手动清空GPU存储
torch.cuda.empty_cache()

实测：提升不大。还没有增加batch_size对速度提升快。待研究...

第四，查询张量基本信息，这在debug深度神经网络时，特别有用

# 查询张量基本信息
a = torch.randn(128,3,32,32)
print(a.type())  #类型
print(a.dim())   #维度
print(a.size())  #尺寸
print(a.shape)   #形状

torch.FloatTensor
4
torch.Size([128, 3, 32, 32])
torch.Size([128, 3, 32, 32])

第五，数据类型转换；Pytorch中，32bit 浮点数的计算效率远高于64bit双精度数

a = torch.randn(128,3,32,32)
a = a.cuda()
print(a.type())  #类型
a = a.cpu()
print(a.type())  #类型
a = a.float()
print(a.type())  #类型
a = a.long()
print(a.type())  #类型

torch.cuda.FloatTensor
torch.FloatTensor
torch.FloatTensor
torch.LongTensor

第六，torch.Tensor与np.ndarray的转换

import numpy as np
a = torch.randn(128,3,32,32)
b = a.numpy()
print(type(b))
c = torch.from_numpy(b)
print(c.type())

<class 'numpy.ndarray'>
torch.FloatTensor

第七，改变Tensor形状，常用于将卷积层特征输入全连接层的情形

# 改变Tensor形状，常用于将卷积层特征输入全连接层的情形
a = torch.randn(128,3,32,32)
b = a.view(-1,16*32*32) # 前级16个卷积单元，展平
print(b.shape)
c = a.reshape(-1, 16*32*32) # reshape可读性更好
print(c.shape)

torch.Size([24, 16384])
torch.Size([24, 16384])

第八，统计代码各部分耗时，找出性能瓶颈

with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
    acc_n = 100
    for epoch in range(2):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to(device), data[1].to(device) # data送到GPU
            # inputs, labels = data
            optimizer.zero_grad() # 将上一次的梯度值清零
            output = net(inputs)  # 前向计算forward()
            loss = criterion(output, labels) # 计算损失值
            loss.backward()       # 反向计算backward()
            running_loss += loss.item() #累积loss值
            optimizer.step()      # 更新神经网络参数

            if i % acc_n == (acc_n-1):
                print('[%d, %5d] loss: %.3f' %
                    (epoch + 1, i + 1, running_loss / acc_n)) #计算平均loss值
                running_loss = 0.0

print('Finished Training!')
print(profile)

profile

或者用命令：

python -m torch.utils.bottleneck xxx.py

python -m torch.utils.bottleneck xxx.py profile报告