Python PyTorch 库详细介绍

import torch# 从列表创建a = torch.tensor([[1, 2], [3, 4]])# 创建特殊张量zeros = torch.zeros(3, 4)          # 全零ones = torch.ones(2, 3)             # 全一rand = torch.randn(3, 3)            # 标准正态分布随机数arange = torch.arange(0, 10, 2)     # [0,2,4,6,8]

if torch.cuda.is_available():    device = torch.device('cuda')    a_gpu = a.to(device)            # 移动到 GPU    a_cpu = a_gpu.cpu()              # 移回 CPU

import numpy as npnp_array = np.array([1, 2, 3])tensor_from_np = torch.from_numpy(np_array)   # NumPy -> Tensornp_from_tensor = tensor_from_np.numpy()       # Tensor -> NumPy

x = torch.tensor([2.0], requires_grad=True)y = x ** 2 + 3 * xy.backward()          # 自动计算梯度 dy/dx = 2*x + 3print(x.grad)         # tensor([7.])

import torch.nn as nnimport torch.nn.functional as Fclass MyModel(nn.Module):    def __init__(self):        super().__init__()        self.conv1 = nn.Conv2d(1, 32, 3)   # 卷积层        self.fc = nn.Linear(32 * 26 * 26, 10)  # 全连接层    def forward(self, x):        x = self.conv1(x)        x = F.relu(x)        # 也可以使用 nn.ReLU()        x = torch.flatten(x, 1)        x = self.fc(x)        return x

optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)# 训练循环中：optimizer.zero_grad()   # 清除之前梯度loss.backward()         # 计算当前梯度optimizer.step()        # 更新参数

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)# 每个 epoch 后调用 scheduler.step()

from torch.utils.data import Datasetclass MyDataset(Dataset):    def __init__(self, data_list, transform=None):        self.data = data_list        self.transform = transform    def __len__(self):        return len(self.data)    def __getitem__(self, idx):        sample = self.data[idx]        if self.transform:            sample = self.transform(sample)        return sample

from torch.utils.data import DataLoaderdataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)for batch in dataloader:    # 训练代码    pass

# 保存torch.save(model.state_dict(), 'model_weights.pth')# 加载model = MyModel()model.load_state_dict(torch.load('model_weights.pth'))

torch.save(model, 'model_full.pth')model = torch.load('model_full.pth')

torch.save({    'epoch': epoch,    'model_state_dict': model.state_dict(),    'optimizer_state_dict': optimizer.state_dict(),    'loss': loss,}, 'checkpoint.pth')

import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms# 1. 数据准备transform = transforms.Compose([    transforms.ToTensor(),    transforms.Normalize((0.1307,), (0.3081,))  # MNIST 均值和标准差])train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)# 2. 定义模型class Net(nn.Module):    def __init__(self):        super().__init__()        self.conv1 = nn.Conv2d(1, 32, 3, 1)        self.conv2 = nn.Conv2d(32, 64, 3, 1)        self.dropout1 = nn.Dropout2d(0.25)        self.fc1 = nn.Linear(9216, 128)        self.dropout2 = nn.Dropout2d(0.5)        self.fc2 = nn.Linear(128, 10)    def forward(self, x):        x = self.conv1(x)        x = nn.ReLU()(x)        x = self.conv2(x)        x = nn.ReLU()(x)        x = nn.MaxPool2d(2)(x)        x = self.dropout1(x)        x = torch.flatten(x, 1)        x = self.fc1(x)        x = nn.ReLU()(x)        x = self.dropout2(x)        x = self.fc2(x)        output = nn.LogSoftmax(dim=1)(x)        return outputdevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')model = Net().to(device)# 3. 定义损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 4. 训练循环epochs = 5for epoch in range(1, epochs + 1):    model.train()    for batch_idx, (data, target) in enumerate(train_loader):        data, target = data.to(device), target.to(device)        optimizer.zero_grad()        output = model(data)        loss = criterion(output, target)        loss.backward()        optimizer.step()        if batch_idx % 100 == 0:            print(f'Epoch {epoch} [{batch_idx*len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}')# 5. 测试model.eval()test_loss = 0correct = 0with torch.no_grad():    for data, target in test_loader:        data, target = data.to(device), target.to(device)        output = model(data)        test_loss += criterion(output, target).item()        pred = output.argmax(dim=1, keepdim=True)        correct += pred.eq(target.view_as(pred)).sum().item()test_loss /= len(test_loader)print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100.*correct/len(test_loader.dataset):.2f}%)')

class MyReLU(torch.autograd.Function):    @staticmethod    def forward(ctx, input):        ctx.save_for_backward(input)        return input.clamp(min=0)    @staticmethod    def backward(ctx, grad_output):        input, = ctx.saved_tensors        grad_input = grad_output.clone()        grad_input[input < 0] = 0        return grad_input

from torch.cuda.amp import autocast, GradScalerscaler = GradScaler()for data, target in train_loader:    optimizer.zero_grad()    withautocast():        output = model(data)        loss = criterion(output, target)    scaler.scale(loss).backward()    scaler.step(optimizer)    scaler.update()

traced_script_module = torch.jit.trace(model, example_input)traced_script_module.save("model.pt")

torch.onnx.export(model, dummy_input, "model.onnx")

model = torch.compile(model)