使用 torch.set_num_threads

 
import torch

# 设置全局线程数
torch.set_num_threads(4)

# 现在导入你的模型和其他 PyTorch 代码...
    

 
注意，torch.set_num_threads 设置的是全局线程数，
它将影响所有的 PyTorch 操作，包括数据加载、模型前向传播等。
    

设置环境变量：

 
import os

# 设置环境变量以限制线程数
os.environ['OMP_NUM_THREADS'] = '4'
os.environ['MKL_NUM_THREADS'] = '4'

# 现在导入 PyTorch 和你的模型
import torch
import torch.nn as nn

# 你的模型定义和加载代码...
    

在数据加载时控制线程数：

 
如果你的瓶颈在于数据加载而不是模型推理，
你可以通过调整 DataLoader 的 num_workers 参数来控制用于数据加载的线程数。
    

 
from torch.utils.data import DataLoader, Dataset

# 假设你有一个自定义的 Dataset 类
dataset = MyDataset()

# 创建 DataLoader 并设置 num_workers
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)
    

torch 调用GPU

device = "cuda:0" if torch.cuda.is_available() else "cpu"
X = X.to(device=device)
y = y.to(device=device)


USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda" if USE_CUDA else "cpu")

 
import torch
use_gpu = torch.cuda.is_available()
if use_gpu:
    model.cuda()

从CPU加载到GPU

 
import torch
x = torch.randn(2, 3)
x
tensor([[ 1.3599,  1.0726, -1.0500],
        [-0.1350, -1.6993, -1.1693]])
    

 
x.device
device(type='cpu')
    

 
这是因为电脑上没有GPU，则默认为CPU 
    

从CPU加载到GPU

 
# Returns a copy of this object in CUDA memory.
x.cuda()

device = "cuda:0" if torch.cuda.is_available() else "cpu"
X = X.to(device=device)

 
a = torch.randn(2, 3, 4)

torch.permute(a, (2, 1, 0))

a.permute(2, 1, 0)

交换两个维度，也只能交换两个维度

 
torch.transpose(a, 0, 1).shape
    

转置

 
import torch
a=torch.tensor([[1,2,3]])
a.shape
torch.Size([1, 3])


b=a.t()
b.shape
torch.Size([3, 1])
    

如果只有2维，可以从行与列的角度去看

shape为（2,3），读作3个维度为2的向量，但这也只是读作，
或者在pytorch看来如此，
或者说对于既定的业务含义有不同的读法，

就数字而言，它也是 2个 维度为3的向量,二者等价
当维度更多时，就是你已经无法一下就说清每个维度的含义，
或者是单纯地从 数学/数字/矩阵 角度去看，
从某个维度取数据时，只改变所取维度，不改变其他维度的大小

 
(2,3) 2行3列，取一行数据，其shape为(3,)
(n1,n2) n1行n2列，取第1维(索引为0)的某个数据，其shape为(n2,)
(n1,n2) n1行n2列，取第2维(索引为1)的某个数据，其shape为(n1,)

(n1,n2,...,nm-1,nm,nm+1,...,nk),
已经无法从行列的角度来说了，但后面的逻辑是一样的

取第nm维后的某个数据，其shape为(n1,n2,...,nm-1,nm+1,...,nk)

pytorch中的维度索引从0开始计数
    

取第2维（索引为1）的某个数

 
```python
import torch

torch.manual_seed(73)
a = torch.ones(1,3,2,2)
print(a)

a = a[:,1,:,:]
print(a.shape)  # torch.Size([1, 2, 2])  这里消失了１维

a = torch.unsqueeze(input=a,dim=1)  # 添加一个维度
print(a.shape) # torch.Size([1, 1, 2, 2])
print(a)

"""
tensor([[[[1., 1.],
          [1., 1.]],

          [[1., 1.],
           [1., 1.]],

          [[1., 1.],
           [1., 1.]]]])
torch.Size([1, 2, 2])
torch.Size([1, 1, 2, 2])
tensor([[[[1., 1.],
          [1., 1.]]]])
"""

```

取某个维度的一个切片

 
a[:,n:m,:]

按维取均值

 
```python 
import torch

torch.manual_seed(73)
a = torch.ones(1,3,2,2)

# 取均值默认会消失一个维度，比如一串数字的均值是一个数值
print(a.mean(dim=1).shape)  # torch.Size([1, 2, 2])
print(a.mean(dim=1,keepdim=True).shape)
print(a)

"""
torch.Size([1, 2, 2])
torch.Size([1, 1, 2, 2])
tensor([[[[1., 1.],
            [1., 1.]],

            [[1., 1.],
            [1., 1.]],

            [[1., 1.],
            [1., 1.]]]])

"""
```

按维取最大值

 
pytorch 取第0维的最大值
values,indices = a.max(dim=0)
values
tensor([4., 5., 6.])

indices
tensor([1, 1, 1])

取第1维的最大值
values,indices = a.max(dim=1)
values
tensor([3., 6.])

unsqueeze，在指定dim上增加一维

 
import torch
probs = torch.tensor([[0.9, 0.1], [0.1, 0.9], [0.8, 0.2]]).unsqueeze(0)
probs
tensor([[[0.9000, 0.1000],
            [0.1000, 0.9000],
            [0.8000, 0.2000]]])

x.view:可兼容onnx转换

 
import torch

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
x_flatten = x.view(2,-1)  # -1 表示自动计算此维度的大小
print(x_flatten)

 
tensor([[1, 2, 3],
        [4, 5, 6]])
     

 
x = torch.randn(3,2,3)
x 
tensor([[[-1.8351, -0.6714, -1.3090],
         [ 1.0558,  1.2539, -0.6256]],

        [[-0.8478,  1.4329, -1.0031],
         [-1.2600,  0.6706,  1.2318]],

        [[ 0.9180,  0.0779, -1.1079],
         [-2.0011, -0.6397, -1.0539]]])

 
x.view(3,-1)
tensor([[-1.8351, -0.6714, -1.3090,  1.0558,  1.2539, -0.6256],
        [-0.8478,  1.4329, -1.0031, -1.2600,  0.6706,  1.2318],
        [ 0.9180,  0.0779, -1.1079, -2.0011, -0.6397, -1.0539]])
        

x.reshape:兼容onnx

 
x.reshape(3,-1)
tensor([[-1.8351, -0.6714, -1.3090,  1.0558,  1.2539, -0.6256],
        [-0.8478,  1.4329, -1.0031, -1.2600,  0.6706,  1.2318],
        [ 0.9180,  0.0779, -1.1079, -2.0011, -0.6397, -1.0539]])

nn.Flatten()：兼容onnx

 
# 展平层
self.flatten = nn.Flatten()

x = self.flatten(x)