当前位置：首页 > news >正文

Pytorch1 PyTorch 官方 QuickStart 超详细笔记｜

news 2026/5/12 21:59:05

适合：刚入门 PyTorch、想快速跑通完整机器学习 pipeline的同学
全文结构：数据准备 → 构建模型 → 训练优化 → 模型评估 → 保存加载 → 单张预测

一、前言

这篇笔记把 PyTorch 官方最快入门案例完整拆解，一行代码一个知识点，帮你快速掌握：

数据集怎么加载
模型怎么定义
训练循环怎么写
模型怎么保存与推理

任务：FashionMNIST 服饰图片分类（10 分类）

二、环境与依赖

importtorchfromtorchimportnnfromtorch.utils.dataimportDataLoaderfromtorchvisionimportdatasetsfromtorchvision.transformsimportToTensor

三、数据准备（Dataset + DataLoader）

1. 下载数据集

# 训练集training_data=datasets.FashionMNIST(root="data",train=True,download=True,transform=ToTensor(),)# 测试集test_data=datasets.FashionMNIST(root="data",train=False,download=True,transform=ToTensor(),)

2. 构建 DataLoader

batch_size=64train_dataloader=DataLoader(training_data,batch_size=batch_size)test_dataloader=DataLoader(test_data,batch_size=batch_size)# 查看数据形状forX,yintest_dataloader:print(f"Shape of X [N, C, H, W]:{X.shape}")print(f"Shape of y:{y.shape}{y.dtype}")break

输出：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28]) Shape of y: torch.Size([64]) torch.int64

四、构建模型（继承 nn.Module）

1. 选择设备

device=torch.accelerator.current_accelerator().typeiftorch.accelerator.is_available()else"cpu"print(f"Using{device}device")

2. 定义网络结构

classNeuralNetwork(nn.Module):def__init__(self):super().__init__()self.flatten=nn.Flatten()self.linear_relu_stack=nn.Sequential(nn.Linear(28*28,512),nn.ReLU(),nn.Linear(512,512),nn.ReLU(),nn.Linear(512,10))defforward(self,x):x=self.flatten(x)logits=self.linear_relu_stack(x)returnlogits model=NeuralNetwork().to(device)print(model)

五、训练相关配置（损失函数 + 优化器）

loss_fn=nn.CrossEntropyLoss()optimizer=torch.optim.SGD(model.parameters(),lr=1e-3)

六、训练函数 & 测试函数

1. 训练函数

deftrain(dataloader,model,loss_fn,optimizer):size=len(dataloader.dataset)model.train()forbatch,(X,y)inenumerate(dataloader):X,y=X.to(device),y.to(device)# 前向传播pred=model(X)loss=loss_fn(pred,y)# 反向传播 + 更新loss.backward()optimizer.step()optimizer.zero_grad()ifbatch%100==0:loss,current=loss.item(),(batch+1)*len(X)print(f"loss:{loss:>7f}[{current:>5d}/{size:>5d}]")

2. 测试函数

deftest(dataloader,model,loss_fn):size=len(dataloader.dataset)num_batches=len(dataloader)model.eval()test_loss,correct=0,0withtorch.no_grad():forX,yindataloader:X,y=X.to(device),y.to(device)pred=model(X)test_loss+=loss_fn(pred,y).item()correct+=(pred.argmax(1)==y).type(torch.float).sum().item()test_loss/=num_batches correct/=sizeprint(f"Test Error:")print(f" Accuracy:{(100*correct):>0.1f}%, Avg loss:{test_loss:>8f}\n")

七、开始训练

epochs=5fortinrange(epochs):print(f"Epoch{t+1}\n-------------------------------")train(train_dataloader,model,loss_fn,optimizer)test(test_dataloader,model,loss_fn)print("Done!")

八、模型保存与加载

1. 保存模型

torch.save(model.state_dict(),"model.pth")print("Saved PyTorch Model State to model.pth")

2. 加载模型

model=NeuralNetwork().to(device)model.load_state_dict(torch.load("model.pth",weights_only=True))

九、单张图片推理

classes=["T-shirt/top","Trouser","Pullover","Dress","Coat","Sandal","Shirt","Sneaker","Bag","Ankle boot",]model.eval()x,y=test_data[0][0],test_data[0][1]withtorch.no_grad():x=x.to(device)pred=model(x)predicted,actual=classes[pred[0].argmax(0)],classes[y]print(f'Predicted: "{predicted}", Actual: "{actual}"')