PyTorch Cheatsheet

torch.tensor([1,2,3])          # 从列表创建
torch.zeros(3,3)               # 3x3零矩阵
torch.ones(2,3,4)              # 2x3x4全1张量
torch.rand(5,5)                # 5x5随机矩阵
torch.arange(0,10,2)           # 类似range()
torch.linspace(0,1,5)          # 线性间隔向量

x.shape                        # 形状
x.dtype                        # 数据类型
x.device                       # 存储设备
x.ndim                         # 维度数
x.numel()                      # 元素总数

x[1:3, :]                      # 行切片
x[:, 2:4]                      # 列切片
x[x > 0.5]                     # 布尔索引

x + y                          # 逐元素加
torch.add(x, y)                # 等价操作
x.mm(y)                        # 矩阵乘法
x @ y                          # Python3.5+矩阵乘法
torch.sqrt(x)                  # 平方根
torch.exp(x)                   # 指数
C1 = A @ B                     # Python3.5+ 语法
C2 = torch.matmul(A, B)        # 等效操作

D = torch.tensor([[1,2],[3,4]])
Dt = D.T                       # tensor([[1,3],
                               #         [2,4]])

x = torch.arange(6)            # tensor([0, 1, 2, 3, 4, 5])
y = x.view(2, 3)               # tensor([[0, 1, 2], 
                               #         [3, 4, 5]])
z = x.reshape(3, 2)            # tensor([[0, 1],
                               #         [2, 3],
                               #         [4, 5]])
E = torch.randn(2,3,4)
Et = E.transpose(1,2)          # 形状变为(2,4,3)

F = torch.randn(2,3)
G = torch.randn(2,5)

H = torch.cat([F,G], dim=1)    # 结果形状(2,8)

x = torch.tensor(2., requires_grad=True)
y = x**2
y.backward()                   # 反向传播
x.grad                         # 梯度值 (4.0)

optimizer.zero_grad()          # 优化器清零
x.grad.data.zero_()           # 手动清零

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

nn.Conv2d(3, 16, kernel_size=3)  # 卷积层
nn.LSTM(100, 256, 2)            # LSTM层
nn.BatchNorm1d(128)             # 批归一化
nn.Dropout(0.5)                 # Dropout层

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __len__(self): return len(data)
    def __getitem__(self, idx): return sample, label

loader = DataLoader(dataset, batch_size=32, shuffle=True)

model.train()
for epoch in range(epochs):
    for data, target in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)                # 模型转移
data = data.to(device)          # 数据转移

torch.save(model.state_dict(), 'model.pth')  # 只保存参数
torch.save(model, 'model_full.pth')          # 保存整个模型

model.load_state_dict(torch.load('model.pth'))
model = torch.load('model_full.pth')

torch.stack([x,y,z])           # 张量堆叠
torch.cat([x,y], dim=0)        # 沿维度连接
torch.chunk(x, 3, dim=1)       # 沿维度分块
torch.split(x, [2,3], dim=1)   # 不均匀分割