PyTorch 学习：从张量到完整训练循环

目录

• 1. PyTorch 是什么
• 2. 环境安装与验证
• 3. 学习路线
• 4. Tensor：PyTorch 的基础数据结构
• 5. Autograd：自动求导
• 6. nn.Module：搭建神经网络
• 7. Dataset 和 DataLoader：组织数据
• 8. 完整示例一：线性回归
• 9. 完整示例二：二分类神经网络
• 10. 完整示例三：FashionMNIST 图像分类
• 11. 保存、加载与推理
• 12. GPU、混合精度与性能优化
• 13. 常见错误与排查
• 14. 建议练习
• 15. 参考资料

1. PyTorch 是什么

PyTorch 是一个用于深度学习和张量计算的 Python 框架。你可以把它理解成：

• NumPy 的增强版：支持 GPU 加速。
• 自动求导系统：可以自动计算梯度。
• 神经网络工具箱：提供 torch.nn、torch.optim、DataLoader 等常用组件。
• 研究和工程都常用的框架：适合快速实验，也能部署模型。

一个典型 PyTorch 训练流程如下：

准备数据 -> 构建模型 -> 定义损失函数 -> 定义优化器 -> 前向传播 -> 计算损失 -> 反向传播 -> 更新参数 -> 保存模型

最常用的模块：

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader

2. 环境安装与验证

建议在项目目录中创建虚拟环境，避免依赖混乱。

2.1 创建虚拟环境

Windows PowerShell：

python –m venv .venv
.\\.venv\\Scripts\\activate
python –m pip install —upgrade pip setuptools wheel

Linux/macOS：

python3 -m venv .venv
source .venv/bin/activate
python -m pip install –upgrade pip setuptools wheel

2.2 安装 PyTorch

CPU 版本通常可以这样安装：

pip install torch torchvision torchaudio

如果你有 NVIDIA GPU，建议去 PyTorch 官方安装页选择你的系统、包管理器和 CUDA 版本，再复制生成的命令：

https://pytorch.org/get-started/locally/

不要随便复制旧教程里的 CUDA 安装命令，因为 PyTorch、CUDA、显卡驱动版本会变化。

2.3 验证安装

新建 check_torch.py：

import torch

print("torch version:", torch.__version__)
print("cuda available:", torch.cuda.is_available())

if torch.cuda.is_available():
print("gpu:", torch.cuda.get_device_name(0))

运行：

python check_torch.py

3. 学习路线

推荐按下面顺序学：

4. Tensor：PyTorch 的基础数据结构

Tensor 是 PyTorch 的核心数据结构，类似 NumPy 数组，但可以放到 GPU 上，也可以参与自动求导。

4.1 创建 Tensor

import torch

a = torch.tensor([1, 2, 3])
b = torch.zeros(2, 3)
c = torch.ones(2, 3)
d = torch.randn(2, 3) # 标准正态分布随机数
e = torch.arange(0, 10, 2) # 0, 2, 4, 6, 8

print(a)
print(b.shape)
print(d.dtype)

4.2 形状、类型、设备

import torch

x = torch.randn(4, 5)

print(x.shape) # torch.Size([4, 5])
print(x.ndim) # 2
print(x.dtype) # torch.float32
print(x.device) # cpu

y = x.to(dtype=torch.float64)
print(y.dtype)

device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)
print(x.device)

4.3 常用操作

import torch

x = torch.arange(12).reshape(3, 4)
print(x)

print(x[0]) # 第 0 行
print(x[:, 1]) # 第 1 列
print(x[1:3, :2]) # 切片

print(x.reshape(2, 6))
print(x.T) # 转置

4.4 数学运算

import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[10.0, 20.0], [30.0, 40.0]])

print(a + b)
print(a * b) # 对应元素相乘
print(a @ b) # 矩阵乘法
print(torch.matmul(a, b))
print(a.mean())
print(a.sum(dim=0)) # 按列求和

4.5 广播机制

广播可以让不同形状的 Tensor 自动扩展后计算。

import torch

x = torch.randn(3, 4)
bias = torch.tensor([1.0, 2.0, 3.0, 4.0])

y = x + bias
print(y.shape) # torch.Size([3, 4])

常见规则：从右往左对齐维度，维度相同或其中一个为 1 时可以广播。

4.6 与 NumPy 互转

import numpy as np
import torch

arr = np.array([1, 2, 3], dtype=np.float32)
x = torch.from_numpy(arr)

y = x.numpy()

print(type(x), x)
print(type(y), y)

注意：CPU Tensor 和 NumPy 数组可能共享内存。修改其中一个，另一个可能也变。

5. Autograd：自动求导

深度学习训练的关键是求梯度。PyTorch 会根据 Tensor 运算动态构建计算图，并通过 backward() 自动反向传播。

5.1 一个最小例子

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1

y.backward()

print("y =", y.item())
print("dy/dx =", x.grad.item()) # 2x + 3，当 x=2 时为 7

5.2 多变量求导

import torch

w = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
x = torch.tensor([4.0, 5.0, 6.0])

y = (w * x).sum()
y.backward()

print(w.grad) # tensor([4., 5., 6.])

5.3 禁用梯度

推理时不需要梯度，使用 torch.no_grad() 可以减少内存和计算开销。

import torch

x = torch.tensor(2.0, requires_grad=True)

with torch.no_grad():
y = x * 10

print(y.requires_grad) # False

5.4 detach

detach() 会创建一个不再追踪梯度的新 Tensor。

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x * x
z = y.detach()

print(y.requires_grad) # True
print(z.requires_grad) # False

6. nn.Module：搭建神经网络

PyTorch 中推荐用 nn.Module 定义模型。

6.1 最简单的模型

import torch
from torch import nn

class LinearModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(in_features=1, out_features=1)

def forward(self, x):
return self.linear(x)

model = LinearModel()
x = torch.tensor([[1.0], [2.0], [3.0]])
y_pred = model(x)

print(y_pred)
print(model)

6.2 Sequential 写法

简单模型也可以用 nn.Sequential：

from torch import nn

model = nn.Sequential(
nn.Linear(10, 32),
nn.ReLU(),
nn.Linear(32, 2)
)

print(model)

6.3 常见层

• nn.Linear：全连接层。
• nn.Conv2d：二维卷积层，常用于图像。
• nn.BatchNorm1d / nn.BatchNorm2d：批归一化。
• nn.ReLU / nn.GELU / nn.Sigmoid：激活函数。
• nn.Dropout：随机丢弃神经元，减少过拟合。
• nn.Embedding：词表索引转向量，常用于 NLP。
• nn.LSTM / nn.GRU / nn.TransformerEncoder：序列建模。

7. Dataset 和 DataLoader：组织数据

实际项目中，不建议把全部训练数据直接写进训练循环。推荐用 Dataset 描述“如何取一条样本”，用 DataLoader 负责批量、打乱和并行加载。

7.1 自定义 Dataset

import torch
from torch.utils.data import Dataset, DataLoader

class ToyDataset(Dataset):
def __init__(self):
self.x = torch.arange(0, 100, dtype=torch.float32).reshape(–1, 1)
self.y = 2 * self.x + 3

def __len__(self):
return len(self.x)

def __getitem__(self, index):
return self.x[index], self.y[index]

dataset = ToyDataset()
loader = DataLoader(dataset, batch_size=8, shuffle=True)

for batch_x, batch_y in loader:
print(batch_x.shape, batch_y.shape)
break

7.2 DataLoader 常用参数

loader = DataLoader(
dataset,
batch_size=32,
shuffle=True,
num_workers=0,
drop_last=False
)

参数说明：

• batch_size：每批样本数。
• shuffle：训练集通常设为 True。
• num_workers：数据加载进程数。Windows 上初学时可以先用 0。
• drop_last：最后一个 batch 不够大小时是否丢弃。

8. 完整示例一：线性回归

目标：让模型学会 y = 3x + 2 + noise。

把下面代码保存为 linear_regression_demo.py 运行。

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def main():
torch.manual_seed(42)

x = torch.linspace(–5, 5, 200).reshape(–1, 1)
noise = torch.randn_like(x) * 0.5
y = 3 * x + 2 + noise

dataset = TensorDataset(x, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for epoch in range(100):
total_loss = 0.0

for batch_x, batch_y in loader:
pred = model(batch_x)
loss = loss_fn(pred, batch_y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

total_loss += loss.item() * batch_x.size(0)

if (epoch + 1) % 10 == 0:
avg_loss = total_loss / len(dataset)
print(f"epoch {epoch + 1:03d}, loss={avg_loss:.4f}")

weight = model.weight.item()
bias = model.bias.item()
print(f"learned: y = {weight:.3f}x + {bias:.3f}")

test_x = torch.tensor([[10.0]])
with torch.no_grad():
test_y = model(test_x)
print("x=10 prediction:", test_y.item())

if __name__ == "__main__":
main()

9. 完整示例二：二分类神经网络

这个例子不用下载数据，直接生成二维点。模型要判断一个点在圆内还是圆外。

保存为 binary_classification_demo.py。

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset, random_split

class MLP(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 32),
nn.ReLU(),
nn.Linear(32, 32),
nn.ReLU(),
nn.Linear(32, 1)
)

def forward(self, x):
return self.net(x)

def accuracy_from_logits(logits, y):
probs = torch.sigmoid(logits)
pred = (probs >= 0.5).float()
return (pred == y).float().mean().item()

def main():
torch.manual_seed(0)

n = 2000
x = torch.rand(n, 2) * 4 – 2
radius = torch.sqrt((x ** 2).sum(dim=1, keepdim=True))
y = (radius < 1.0).float()

dataset = TensorDataset(x, y)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) – train_size
train_set, val_set = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256, shuffle=False)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = MLP().to(device)
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(30):
model.train()
train_loss = 0.0

for batch_x, batch_y in train_loader:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)

logits = model(batch_x)
loss = loss_fn(logits, batch_y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

train_loss += loss.item() * batch_x.size(0)

model.eval()
val_acc_sum = 0.0
val_count = 0

with torch.no_grad():
for batch_x, batch_y in val_loader:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)

logits = model(batch_x)
batch_acc = accuracy_from_logits(logits, batch_y)

val_acc_sum += batch_acc * batch_x.size(0)
val_count += batch_x.size(0)

avg_loss = train_loss / len(train_set)
val_acc = val_acc_sum / val_count
print(f"epoch {epoch + 1:02d}, loss={avg_loss:.4f}, val_acc={val_acc:.4f}")

sample = torch.tensor([[0.2, 0.3], [1.5, 1.5]], device=device)
model.eval()
with torch.no_grad():
probs = torch.sigmoid(model(sample))
print("sample probabilities:", probs.cpu().squeeze().tolist())

if __name__ == "__main__":
main()

关键点：

• 二分类推荐输出一个 logit，然后用 nn.BCEWithLogitsLoss()。
• 不要在模型最后手动加 Sigmoid 再送进 BCEWithLogitsLoss，它内部已经处理了数值稳定问题。
• 训练时用 model.train()，验证和推理时用 model.eval()。

10. 完整示例三：FashionMNIST 图像分类

这个例子需要 torchvision，第一次运行会下载 FashionMNIST 数据集。

保存为 fashion_mnist_demo.py。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class SmallCNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(16, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32 * 7 * 7, 128),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(128, 10)
)

def forward(self, x):
x = self.features(x)
return self.classifier(x)

def evaluate(model, loader, device):
model.eval()
correct = 0
total = 0

with torch.no_grad():
for x, y in loader:
x = x.to(device)
y = y.to(device)
logits = model(x)
pred = logits.argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)

return correct / total

def main():
torch.manual_seed(42)

transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])

train_set = datasets.FashionMNIST(
root="./data",
train=True,
download=True,
transform=transform
)
test_set = datasets.FashionMNIST(
root="./data",
train=False,
download=True,
transform=transform
)

train_loader = DataLoader(train_set, batch_size=64, shuffle=True, num_workers=0)
test_loader = DataLoader(test_set, batch_size=256, shuffle=False, num_workers=0)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SmallCNN().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(5):
model.train()
total_loss = 0.0

for x, y in train_loader:
x = x.to(device)
y = y.to(device)

logits = model(x)
loss = loss_fn(logits, y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

total_loss += loss.item() * x.size(0)

train_loss = total_loss / len(train_set)
test_acc = evaluate(model, test_loader, device)
print(f"epoch {epoch + 1}, train_loss={train_loss:.4f}, test_acc={test_acc:.4f}")

torch.save(model.state_dict(), "fashion_mnist_cnn.pt")
print("saved to fashion_mnist_cnn.pt")

if __name__ == "__main__":
main()

多分类任务常见组合：

model output: shape = [batch_size, num_classes]
target: shape = [batch_size]，类别编号如 0、1、2
loss_fn = nn.CrossEntropyLoss()

注意：使用 CrossEntropyLoss 时，模型最后一层不要加 Softmax，因为损失函数内部会处理。

11. 保存、加载与推理

推荐保存模型参数 state_dict，而不是直接保存整个模型对象。

11.1 保存

torch.save(model.state_dict(), "model.pt")

11.2 加载

加载前需要先创建相同结构的模型。

import torch
from torch import nn

class SmallModel(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(4, 16),
nn.ReLU(),
nn.Linear(16, 3)
)

def forward(self, x):
return self.net(x)

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

model = SmallModel().to(device)
model.load_state_dict(torch.load("model.pt", map_location=device))
model.eval()

x = torch.randn(1, 4, device=device)
with torch.no_grad():
logits = model(x)
pred = logits.argmax(dim=1)

print(pred.item())

11.3 保存更多训练信息

如果希望中断后继续训练，可以保存 checkpoint。

checkpoint = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": epoch,
"loss": loss.item()
}

torch.save(checkpoint, "checkpoint.pt")

加载：

checkpoint = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1

12. GPU、混合精度与性能优化

12.1 使用 device

模型和数据必须在同一个设备上。

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

model = model.to(device)
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)

如果报错：

Expected all tensors to be on the same device

通常是模型在 GPU，但数据还在 CPU，或反过来。

12.2 混合精度训练

混合精度常用于 NVIDIA GPU，可以减少显存占用并提升速度。

import torch

scaler = torch.amp.GradScaler("cuda")

for batch_x, batch_y in train_loader:
batch_x = batch_x.to("cuda")
batch_y = batch_y.to("cuda")

optimizer.zero_grad()

with torch.amp.autocast(device_type="cuda"):
logits = model(batch_x)
loss = loss_fn(logits, batch_y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

如果没有 CUDA，就先不要使用这段代码。

12.3 torch.compile

PyTorch 2.x 提供了 torch.compile，可以尝试编译模型以提升运行速度。

model = torch.compile(model)

建议：

• 先保证普通训练代码正确。
• 再尝试 torch.compile(model)。
• 如果遇到兼容问题，先移除它排查。

13. 常见错误与排查

13.1 形状不匹配

错误示例：

RuntimeError: mat1 and mat2 shapes cannot be multiplied

排查方式：

print(x.shape)
print(logits.shape)
print(y.shape)

尤其注意 nn.Linear(in_features, out_features) 的 in_features 是否等于输入最后一维。

13.2 标签类型不对

CrossEntropyLoss 要求：

logits.dtype == torch.float32
y.dtype == torch.long

示例：

y = y.long()
loss = nn.CrossEntropyLoss()(logits, y)

13.3 忘记清梯度

PyTorch 默认会累积梯度，所以训练时必须清空旧梯度：

optimizer.zero_grad()
loss.backward()
optimizer.step()

13.4 训练和推理模式混用

带有 Dropout、BatchNorm 的模型，训练和推理行为不同。

model.train() # 训练
model.eval() # 验证、测试、推理

13.5 把概率传给 CrossEntropyLoss

错误写法：

prob = torch.softmax(logits, dim=1)
loss = nn.CrossEntropyLoss()(prob, y)

正确写法：

loss = nn.CrossEntropyLoss()(logits, y)

14. 建议练习

练习 1：Tensor 基础

写代码完成：

• 创建一个形状为 [5, 3] 的随机 Tensor。
• 取第 2 行。
• 计算每一列的平均值。
• 转成 GPU Tensor，如果没有 GPU 就保持 CPU。

参考答案：

import torch

x = torch.randn(5, 3)
print(x[1])
print(x.mean(dim=0))

device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)
print(x.device)

练习 2：改造线性回归

把线性回归示例中的真实函数改成：

y = –4 * x + 7 + noise

观察模型最后学到的 weight 和 bias 是否接近 -4 和 7。

练习 3：给二分类模型加层

把二分类示例中的模型改成：

self.net = nn.Sequential(
nn.Linear(2, 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, 1)
)

比较训练速度和验证准确率。

练习 4：保存并加载模型

在二分类例子训练结束后保存：

torch.save(model.state_dict(), "circle_classifier.pt")

然后新建模型，加载参数，对几个点做推理。

练习 5：手写一个 Dataset

写一个 CSVDataset，从 CSV 文件读取特征和标签。可以先假设最后一列是标签，其余列是特征。

提示：

import pandas as pd
import torch
from torch.utils.data import Dataset

class CSVDataset(Dataset):
def __init__(self, path):
df = pd.read_csv(path)
self.x = torch.tensor(df.iloc[:, :–1].values, dtype=torch.float32)
self.y = torch.tensor(df.iloc[:, –1].values, dtype=torch.long)

def __len__(self):
return len(self.x)

def __getitem__(self, index):
return self.x[index], self.y[index]

15. 参考资料

• PyTorch 官方安装页：https://pytorch.org/get-started/locally/
• PyTorch 官方教程首页：https://docs.pytorch.org/tutorials/
• PyTorch Quickstart：https://docs.pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html
• Tensor 教程：https://docs.pytorch.org/tutorials/beginner/basics/tensorqs_tutorial.html
• Autograd 教程：https://docs.pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html

16. 面试常见问题

16.1 什么是 PyTorch 的动态图机制？

PyTorch 默认采用动态图，计算图是在前向传播时动态构建的。它的好处是调试方便、代码直观，适合研究和快速迭代。

16.2 requires_grad 是什么？

它表示这个 Tensor 是否需要记录梯度。参数一般需要梯度，输入数据通常不需要。

16.3 backward() 做了什么？

它会根据当前标量损失，沿计算图自动反向传播，计算叶子节点的梯度并写入 tensor.grad。

16.4 为什么训练前要 optimizer.zero_grad()？

因为 PyTorch 默认会累积梯度，如果不清零，上一次和这一次的梯度会叠加，导致参数更新错误。

16.5 model.train() 和 model.eval() 有什么区别？

train() 启用训练行为，eval() 启用推理行为。Dropout、BatchNorm 在这两种模式下行为不同。

16.6 为什么 CrossEntropyLoss 前面不要再加 Softmax？

因为 CrossEntropyLoss 内部已经包含了 LogSoftmax 和负对数似然计算，直接传 logits 更稳定。

16.7 二分类为什么常用 BCEWithLogitsLoss？

它把 Sigmoid 和二元交叉熵合在一起，数值更稳定，适合输出单个 logit 的二分类任务。

16.8 state_dict 和直接保存整个模型有什么区别？

state_dict 只保存参数，轻量、兼容性更好，是推荐方式。直接保存整个模型对象更依赖代码结构，迁移性较差。

16.9 PyTorch 中 CPU 和 GPU Tensor 可以直接混用吗？

不可以。模型和数据必须在同一个设备上，否则会报 device 不一致错误。

16.10 DataLoader 的作用是什么？

它负责把 Dataset 组织成 batch，并支持打乱、并行加载和最后一批处理。

16.11 torch.no_grad() 适合什么时候用？

验证和推理时用。它会关闭梯度记录，减少内存占用并提升速度。

16.12 detach() 有什么用？

它会把 Tensor 从当前计算图中分离出来，后续运算不再追踪梯度。

16.13 nn.Module 为什么要继承？

这样可以统一管理参数、子模块、前向逻辑和设备迁移，是 PyTorch 组织模型的标准方式。

16.14 训练时 loss 不下降怎么办？

优先检查这些点：

• 学习率是否过大或过小。
• 标签类型和损失函数是否匹配。
• 输入输出维度是否正确。
• 模型是否忘记 train()。
• 梯度是否正常反传。
• 数据是否有问题，比如归一化、噪声、标签错误。

16.15 面试时如何概括 PyTorch 的核心流程？

可以这样回答：

先用 Dataset/DataLoader 组织数据，再用 nn.Module 定义模型，选择合适的 loss 和 optimizer，
训练时做前向传播、算损失、清梯度、反向传播、参数更新，验证和推理时切换到 eval 模式，
最后保存 state_dict 方便部署和复现。

训练循环固定模式：

pred = model(batch_x)
loss = loss_fn(pred, batch_y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

这四步是 PyTorch 入门最重要的肌肉记忆。

附录：最小训练模板

以后写 PyTorch 项目，可以从这个模板开始改。

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

class Model(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(10, 32),
nn.ReLU(),
nn.Linear(32, 2)
)

def forward(self, x):
return self.net(x)

def main():
torch.manual_seed(42)

x = torch.randn(1000, 10)
y = torch.randint(0, 2, (1000,))

dataset = TensorDataset(x, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = Model().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
model.train()
total_loss = 0.0

for batch_x, batch_y in loader:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)

logits = model(batch_x)
loss = loss_fn(logits, batch_y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

total_loss += loss.item() * batch_x.size(0)

print(f"epoch={epoch + 1}, loss={total_loss / len(dataset):.4f}")

torch.save(model.state_dict(), "model.pt")

if __name__ == "__main__":
main()