张量与自动微分：深度学习的计算地基

"""张量的运算、索引、变形与拼接：把数据摆成网络要的形状。"""
import torch

torch.manual_seed(0)

# ---------- 一、逐元素运算：形状相同，对应位置各算各的 ----------
a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5, 6], [7, 8]])
print(a.add(10))       # 返回新张量，a 本身不变
print(a.add_(10))      # 带下划线 = 原地改，a 被改掉了
print(torch.mul(a, b)) # 点乘（Hadamard），等价写法 a * b

# ---------- 二、矩阵乘法：形状必须接得上 (n,m) @ (m,p) -> (n,p) ----------
m1 = torch.tensor([[1, 2], [3, 4], [5, 6]])   # (3, 2)
m2 = torch.tensor([[5, 6], [7, 8]])           # (2, 2)
print(m1 @ m2)                 # (3, 2)，运算符写法
print(torch.matmul(m1, m2))    # 等价，且支持带 batch 的高维张量

# ---------- 三、常用统计函数：注意 dim ----------
d = torch.randint(0, 10, [2, 3], dtype=torch.float64)  # mean 要求浮点类型
print(d.mean())          # 全局均值
print(d.mean(dim=0))     # 沿 0 轴压扁 -> 每列一个均值
print(d.mean(dim=1))     # 沿 1 轴压扁 -> 每行一个均值
print(d.sum(), d.sqrt(), d.exp(), d.log())

# ---------- 四、索引：取子块的四种写法 ----------
g = torch.randint(0, 10, [4, 5])
print(g[0])                     # 第 0 行
print(g[:, 0])                  # 第 0 列
print(g[[0, 1], [1, 2]])        # 花式索引：(0,1) 和 (1,2) 两个点
print(g[:3, :2])                # 范围索引：前 3 行的前 2 列
print(g[g[:, 2] > 5])           # 布尔索引：第 2 列大于 5 的那些行

# ---------- 五、变形：网络报形状错误时最常动的就是这几个 ----------
t = torch.tensor([[10, 20, 30], [40, 50, 60]])   # (2, 3)
print(t.reshape(1, 6).shape)                      # (1, 6)，最通用
print(t.unsqueeze(dim=0).shape)                   # (1, 2, 3)，补一个维度
print(t.unsqueeze(0).squeeze().shape)             # (2, 3)，去掉长度为 1 的维度

h = torch.randint(0, 10, [3, 4, 5])
print(torch.transpose(h, 1, 2).shape)             # (3, 5, 4)，交换两个维度
print(h.permute([1, 2, 0]).shape)                 # (4, 5, 3)，一次换多个

# view 只认内存连续的张量，transpose 之后必须先 contiguous
swapped = torch.transpose(t, 0, 1)
print(swapped.is_contiguous())                    # False
print(swapped.contiguous().view(2, 3).shape)      # (2, 3)

# ---------- 六、拼接：cat 不加维度，stack 加一个新维度 ----------
p = torch.randint(0, 10, [1, 2, 3])
q = torch.randint(0, 10, [1, 2, 3])
print(torch.cat([p, q], dim=0).shape)      # (2, 2, 3)，在已有维度上接长
print(torch.stack([p, q], dim=0).shape)    # (2, 1, 2, 3)，多出一维

"""链式法则实证：三层复合函数，手算一遍再让 PyTorch 算一遍，对上才算懂。

复合关系：
    u = 3 * w + 2
    v = u ** 2
    loss = 5 * v

手推：
    dloss/dv = 5
    dv/du    = 2u
    du/dw    = 3
    dloss/dw = 5 * 2u * 3 = 30u = 30 * (3w + 2)
w = 1 时，u = 5，dloss/dw = 150。
"""
import torch


def main():
    w = torch.tensor(1.0, requires_grad=True)

    u = 3 * w + 2
    v = u ** 2
    loss = 5 * v

    # 每一步都带着 grad_fn，这就是反向时要走的路
    print('u    =', u.item(), ' grad_fn =', u.grad_fn)
    print('v    =', v.item(), ' grad_fn =', v.grad_fn)
    print('loss =', loss.item(), ' grad_fn =', loss.grad_fn)

    loss.backward()
    print('PyTorch 算出的 dloss/dw =', w.grad.item())
    print('手推公式 30*(3w+2) 在 w=1 处 =', 30 * (3 * 1 + 2))

    # ------------------------------------------------------------ 梯度会累加
    # 不清零就再反向一次，梯度变成两倍，这就是忘了 zero_grad() 的后果
    w2 = torch.tensor(1.0, requires_grad=True)
    for i in range(3):
        y = 5 * (3 * w2 + 2) ** 2
        y.backward()
        print('第 %d 次 backward 之后 w2.grad = %.1f' % (i + 1, w2.grad.item()))
    # 输出 150 / 300 / 450：每次都往上叠，没有任何报错

    # ------------------------------------------------------------ 正确做法
    w3 = torch.tensor(1.0, requires_grad=True)
    for i in range(3):
        y = 5 * (3 * w3 + 2) ** 2
        if w3.grad is not None:
            w3.grad.zero_()        # 清零之后每次都是干净的 150
        y.backward()
        print('清零后第 %d 次 w3.grad = %.1f' % (i + 1, w3.grad.item()))

    # ------------------------------------------------------------ 中间结果没有梯度
    print('中间张量 v 的 grad:', v.grad)   # None，并伴随一条 UserWarning
    # 确实需要时要显式保留：
    w4 = torch.tensor(1.0, requires_grad=True)
    u4 = 3 * w4 + 2
    u4.retain_grad()
    (5 * u4 ** 2).backward()
    print('retain_grad 之后 u4.grad =', u4.grad.item())   # 10u = 50


if __name__ == '__main__':
    main()

"""广播（broadcasting）：形状不同的张量为什么有时能直接相加。

广播是把小的那个张量"复制"到大的形状上，但它不真的复制内存。
好处是省代码，坏处是**形状写错时它不报错**，安安静静算出一个错的结果。
训练里最隐蔽的一类 bug 就来自这里。
"""
import torch

# ---------------------------------------------------------------- 一、广播的两条规则
# 从最后一个维度往前对齐，逐维比较：
#   1) 长度相等            -> 直接对上
#   2) 其中一个长度是 1    -> 把这一维拉伸成另一个的长度
# 任何一维两条都不满足，才会报错。

a = torch.ones(3, 4)            # (3, 4)
b = torch.tensor([1., 2., 3., 4.])   # (4,)  -> 前面补 1 变 (1, 4) -> 拉伸成 (3, 4)
print((a + b).shape)            # (3, 4)

c = torch.ones(3, 1)            # (3, 1) -> 第 1 维拉伸成 4
print((a + c).shape)            # (3, 4)

d = torch.ones(3, 5)
try:
    _ = a + d                   # 4 和 5 既不相等、也都不是 1
except RuntimeError as e:
    print('形状不兼容:', e)


# ---------------------------------------------------------------- 二、真实会坑到人的场景
# 模型输出 (batch, 1)，标签 (batch,)，两者算 MSE
y_pred = torch.randn(8, 1)      # 网络输出天然带最后一维
y_true = torch.randn(8)         # 从 DataLoader 里直接拿到的标签

wrong = ((y_pred - y_true) ** 2).mean()     # 广播成 (8, 8)，算了 64 个差值
right = ((y_pred - y_true.reshape(-1, 1)) ** 2).mean()   # 只算 8 个
print('形状被悄悄广播成:', (y_pred - y_true).shape)
print('错的损失 %.6f，对的损失 %.6f' % (wrong.item(), right.item()))
# 两个数都是"正常"的浮点数，不看形状根本发现不了


# ---------------------------------------------------------------- 三、自检习惯
# 算损失前把两边形状打出来，一秒钟的事，能省掉半天的排查
print('y_pred', y_pred.shape, 'y_true', y_true.shape)
assert y_pred.shape == y_true.reshape(-1, 1).shape

"""把张量搬到 GPU：三种写法，以及一条必须记住的规矩。

规矩：参与同一次运算的张量必须在同一个设备上。
模型在 GPU、数据在 CPU，会直接抛 RuntimeError，
报错原文是 Expected all tensors to be on the same device。
"""
import torch


# ---------------------------------------------------------------- 先探测环境
# 标准写法：有显卡就用显卡，没有就退回 CPU，脚本两边都能跑
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('当前设备:', device)
print('可用显卡数量:', torch.cuda.device_count())


# ---------------------------------------------------------------- 写法一：建的时候就指定
a = torch.tensor([10, 20, 30], device=device)
print(a.device)


# ---------------------------------------------------------------- 写法二：.cuda() / .cpu()
b = torch.tensor([10, 20, 30])
if torch.cuda.is_available():
    b = b.cuda()          # 搬到默认显卡
    print(b.device)
    b = b.cpu()           # 再搬回来
    print(b.device)


# ---------------------------------------------------------------- 写法三：.to(device)，推荐
# 只有这一种能让同一份代码在有卡/没卡的机器上都不用改
c = torch.tensor([10, 20, 30]).to(device)
print(c.device)


# ---------------------------------------------------------------- 跨设备运算会报错
cpu_tensor = torch.tensor([1., 2., 3.])
if torch.cuda.is_available():
    gpu_tensor = cpu_tensor.to('cuda:0')
    try:
        _ = cpu_tensor + gpu_tensor      # 一个在内存、一个在显存
    except RuntimeError as e:
        print('跨设备相加报错:', e)
    # 正确做法：先统一设备
    print((cpu_tensor.to(device) + gpu_tensor).device)


# ---------------------------------------------------------------- 训练脚本里的固定套路
# model = MyModel().to(device)
# for bx, by in dataloader:
#     bx, by = bx.to(device), by.to(device)   # 每个 batch 都要搬
#     ...
# 注意：模型搬一次就够（参数常驻显存），数据要每个 batch 搬一次。


# ---------------------------------------------------------------- 取回 CPU 才能转 NumPy
d = torch.tensor([1., 2., 3.], requires_grad=True).to(device)
# .detach() 断开计算图，.cpu() 搬回内存，两步都不能少
print(d.detach().cpu().numpy())

"""张量的创建、类型与形状：最小可跑例子。"""
import torch

# ① 用现成数据建张量：数据是什么类型，张量就跟着是什么类型
scalar = torch.tensor(10)                      # 0 维，标量
vector = torch.tensor([10., 20., 30.])         # 1 维，向量，float32
matrix = torch.tensor([[1, 2, 3], [4, 5, 6]])  # 2 维，矩阵，int64
print(scalar.shape, vector.shape, matrix.shape)   # torch.Size([]) ([3]) ([2, 3])

# ② 只给形状建张量：内存里是什么就是什么，是脏数据，必须自己覆盖掉
dirty = torch.Tensor(2, 3)      # 大写 T，按形状开空间，值不确定
print(dirty.dtype)              # torch.float32

# ③ 建规则张量：真实训练里用得最多
print(torch.zeros(2, 3))                 # 全 0
print(torch.ones(2, 3))                  # 全 1
print(torch.full([2, 3], 10))            # 全为指定值
print(torch.arange(0, 10, 2))            # 按步长取点：0 2 4 6 8
print(torch.linspace(0, 9, 10))          # 按个数等分

# ④ 随机张量：做实验前先固定种子，否则每次结果都不一样，没法复现
torch.manual_seed(100)
print(torch.randn(2, 3))                 # 标准正态分布
print(torch.randint(0, 10, [2, 3]))      # 指定区间的随机整数

# ⑤ 类型转换：两套写法等价，挑顺手的
x = torch.full([2, 3], 10)
print(x.dtype)                           # torch.int64
print(x.type(torch.DoubleTensor).dtype)  # torch.float64
print(x.double().dtype)                  # torch.float64

# ⑥ 和 NumPy 互转：默认共享内存，改一个另一个跟着变
import numpy as np
t = torch.tensor([2, 3, 4])
n = t.numpy()            # 共享内存
n[0] = 100
print(t)                 # tensor([100, 3, 4])，被连带改了
n2 = t.numpy().copy()    # 加 copy 才断开
n2[0] = 999
print(t)                 # tensor([100, 3, 4])，这次没被影响

# ⑦ 单元素张量取出 Python 数字：训练循环里累加 loss 全靠它
print(torch.tensor([30.]).item())        # 30.0

"""自动微分最小例子：一次前向、一次反向、读一次梯度。"""
import torch

# requires_grad=True 是自动微分的总开关，且必须是浮点类型
x = torch.tensor(10, requires_grad=True, dtype=torch.float32)

y = 2 * x ** 2          # 前向：PyTorch 边算边把这条链记下来
print(y)                # tensor(200., grad_fn=<MulBackward0>)，grad_fn 就是轨迹
print(y.grad_fn)        # 记录了这一步是乘法算出来的

y.backward()            # 反向：沿轨迹倒推，只对标量成立
print(x.grad)           # tensor(40.)  —— 手算：y' = 4x，代 x=10 得 40

# 如果 y 是向量，backward 会报错，先 sum() 变标量再求
v = torch.tensor([10., 20.], requires_grad=True)
(2 * v ** 2).sum().backward()
print(v.grad)           # tensor([40., 80.])  —— 4x 分别代 10 和 20

"""Dataset 与 DataLoader：数据怎么成批进网络。

两种写法：
  ① 数据已经在内存里 -> TensorDataset 包一下就行
  ② 数据在磁盘/数据库/需要按需加载 -> 自定义 Dataset，写三个方法
"""
import torch
from torch.utils.data import Dataset, TensorDataset, DataLoader


# ---------------------------------------------------------------- ① 最省事的写法
x = torch.randn(103, 4)                      # 103 条样本，故意不是 batch 的整数倍
y = torch.randint(0, 3, (103,))              # 三分类，标签必须是 int64

ds = TensorDataset(x, y)
print('样本总数:', len(ds))
print('第 0 条:', ds[0][0].shape, ds[0][1].item())

loader = DataLoader(ds, batch_size=16, shuffle=True, drop_last=False)
for i, (bx, by) in enumerate(loader):
    print('第 %d 批  x=%s  y=%s' % (i, tuple(bx.shape), tuple(by.shape)))
# 一共 7 批：前 6 批各 16 条，最后一批只有 7 条
# drop_last=True 会把这不满的一批直接丢掉，训练时可选，验证时别丢


# ---------------------------------------------------------------- ② 自定义 Dataset
class MyDataset(Dataset):
    """三个方法是硬性接口，少一个 DataLoader 就用不了。"""

    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        """告诉 DataLoader 一共有多少条。"""
        return len(self.features)

    def __getitem__(self, index):
        """按下标返回一条样本。真实项目里这里才去读图片、查数据库。

        返回的必须是张量，且每条的形状要一致，否则没法拼成一个 batch。
        """
        return self.features[index], self.labels[index]


custom = MyDataset(x, y)
loader2 = DataLoader(custom, batch_size=32, shuffle=False)
bx, by = next(iter(loader2))
print('自定义 Dataset 的一批:', bx.shape, by.shape)


# ---------------------------------------------------------------- 为什么必须 shuffle
# 如果数据按标签排好序（前 1/3 全是类别 0，中间全是 1……）
sorted_y = torch.cat([torch.zeros(34), torch.ones(34), torch.full((35,), 2.)]).long()
bad = DataLoader(TensorDataset(x, sorted_y), batch_size=16, shuffle=False)
first_batch_labels = next(iter(bad))[1]
print('不打乱时第一批的标签:', first_batch_labels.unique().tolist())
# 整批只有一个类别，这一步算出来的梯度方向严重偏向该类，
# 模型会先被拽向"全猜 0"，下一批又被拽向"全猜 1"，训练来回震荡。

"""用自动微分手写梯度下降，求 y = x**2 + 20 的极小值点。

这一份把训练循环剥到最干净：没有 model，没有 optimizer，
只剩「前向 → 清零 → 反向 → 更新」四步，看清楚每一步在动谁。
"""
import torch


def main():
    # 1. 起点随便选一个远离最优解的位置
    x = torch.tensor(10, requires_grad=True, dtype=torch.float32)
    lr = 0.01          # 学习率：每次沿梯度反方向走多远
    epochs = 1000

    print('起点   x=%.6f  y=%.6f' % (x.item(), (x ** 2 + 20).item()))

    for i in range(1, epochs + 1):
        # 1) 前向计算：每轮都要重新算，因为 x 变了
        y = x ** 2 + 20

        # 2) 梯度清零：.grad 默认累加，不清零会把历史梯度一起算进来
        #    第一轮 x.grad 还不存在，所以要判空
        if x.grad is not None:
            x.grad.zero_()

        # 3) 反向传播：把 dy/dx 算出来写进 x.grad，此时 x 本身还没动
        y.backward()

        # 4) 参数更新：必须走 x.data，直接写 x = x - lr * x.grad
        #    会把 x 换成一个新的、不再是叶子节点的张量，下一轮就没有 .grad 了
        x.data = x.data - lr * x.grad

        if i % 200 == 0 or i == 1:
            print('第 %4d 轮  x=%.6f  y=%.6f  梯度=%.6f'
                  % (i, x.item(), (x ** 2 + 20).item(), x.grad.item()))

    # 理论最优解：y' = 2x = 0 -> x = 0，y = 20
    print('收敛到 x=%.8f，最小值 y=%.6f（理论值 x=0, y=20）'
          % (x.item(), (x ** 2 + 20).item()))


if __name__ == '__main__':
    main()

"""线性回归完整案例：第一次把 PyTorch 的四件套凑齐。

数据集 -> DataLoader -> nn.Linear -> MSELoss + SGD -> 训练循环
后面每个模型都是这套骨架的放大版，只是层数和损失函数换了。
"""
import torch
from torch import nn, optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.datasets import make_regression


def create_dataset():
    """造一批带噪声的线性数据：y ≈ coef * x + 14.5。"""
    x, y, coef = make_regression(n_samples=100, n_features=1, noise=10,
                                 coef=True, bias=14.5, random_state=0)
    # sklearn 给的是 numpy，转成张量；统一成 float32，否则 nn.Linear 会报 dtype 不匹配
    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)
    return x, y, coef


def train(x, y, epochs=100, lr=1e-2, batch_size=16):
    # 1) 数据：TensorDataset 负责按下标取一条，DataLoader 负责成批、打乱
    dataloader = DataLoader(TensorDataset(x, y), batch_size=batch_size, shuffle=True)

    # 2) 模型：一个输入特征、一个输出值，内部就是 y = w*x + b
    model = nn.Linear(in_features=1, out_features=1)

    # 3) 损失函数与优化器：优化器要拿到 model.parameters() 才知道该更新谁
    criterion = nn.MSELoss()
    optimizer = optim.SGD(params=model.parameters(), lr=lr)

    epoch_loss = []
    for _ in range(epochs):
        total_loss, batch_num = 0.0, 0
        for train_x, train_y in dataloader:
            # 前向：模型吃 (batch, 1)，吐 (batch, 1)
            y_pred = model(train_x)
            # 真实值是 (batch,)，要 reshape 成 (batch, 1) 才能和预测值对齐
            loss = criterion(y_pred, train_y.reshape(-1, 1))

            optimizer.zero_grad()   # 清零
            loss.backward()         # 反向
            optimizer.step()        # 更新

            total_loss += loss.item()
            batch_num += 1
        epoch_loss.append(total_loss / batch_num)

    return model, epoch_loss


def main():
    x, y, coef = create_dataset()
    model, epoch_loss = train(x, y)

    # 训练出来的 w 应该逼近 coef，b 应该逼近 14.5
    print('真实斜率 coef = %.4f，真实截距 = 14.5' % float(coef))
    print('训练得到 w = %.4f，b = %.4f'
          % (model.weight.item(), model.bias.item()))
    print('首轮平均损失 %.4f -> 末轮平均损失 %.4f'
          % (epoch_loss[0], epoch_loss[-1]))


if __name__ == '__main__':
    main()

"""线性回归的"裸手"版：不用 nn.Linear、不用 MSELoss、不用 optim。

和 linear_regression.py 对照着看，就能看清 PyTorch 的四件套
分别替你做掉了哪一件事。这一份里每一行都是你自己写的。
"""
import torch


def create_dataset(n=100, true_w=8.0, true_b=14.5, seed=0):
    """造 y = true_w * x + true_b + 噪声 的一维数据。"""
    torch.manual_seed(seed)
    x = torch.randn(n, 1)                      # (n, 1)
    noise = torch.randn(n, 1) * 0.5
    y = true_w * x + true_b + noise            # (n, 1)
    return x, y


def data_iter(x, y, batch_size):
    """手写的 DataLoader：打乱下标，按批切片吐出去。"""
    n = len(x)
    idx = torch.randperm(n)                    # 随机排列，等价于 shuffle=True
    for start in range(0, n, batch_size):
        batch_idx = idx[start:start + batch_size]
        yield x[batch_idx], y[batch_idx]


def main():
    x, y = create_dataset()

    # 参数自己建，自己打开自动微分开关
    # nn.Linear 替你做的就是这两行 + forward
    w = torch.zeros(1, 1, requires_grad=True)
    b = torch.zeros(1, requires_grad=True)

    lr, epochs, batch_size = 0.05, 60, 16

    for epoch in range(epochs):
        total, cnt = 0.0, 0
        for bx, by in data_iter(x, y, batch_size):
            # 1) 前向：nn.Linear 内部就是这一行矩阵乘法加偏置
            y_pred = bx @ w + b

            # 2) 损失：MSELoss 内部就是这一行
            loss = ((y_pred - by) ** 2).mean()

            # 3) 清零：optimizer.zero_grad() 替你遍历所有参数做这件事
            if w.grad is not None:
                w.grad.zero_()
                b.grad.zero_()

            # 4) 反向：这一步谁也替不了，autograd 的核心
            loss.backward()

            # 5) 更新：optimizer.step() 内部就是这两行
            with torch.no_grad():              # 不让更新动作本身进计算图
                w -= lr * w.grad
                b -= lr * b.grad

            total += loss.item() * len(by)
            cnt += len(by)

        if (epoch + 1) % 20 == 0:
            print('epoch %3d  loss %.6f  w=%.4f  b=%.4f'
                  % (epoch + 1, total / cnt, w.item(), b.item()))

    print('真实值 w=8.0000  b=14.5000')
    print('学到值 w=%.4f  b=%.4f' % (w.item(), b.item()))


if __name__ == '__main__':
    main()

"""训练脚本骨架模板：把 TODO 填掉就是一个能跑的回归/分类训练。

复制这份文件改，比从空文件开始写少踩一半的坑。
"""
import torch
from torch import nn, optim
from torch.utils.data import TensorDataset, DataLoader


# ---------------------------------------------------------------- 数据
def create_dataset():
    # TODO: 换成你自己的数据来源（read_csv / numpy / 现成 Dataset）
    x = torch.randn(200, 4)                  # TODO: 特征，形状 (样本数, 特征数)
    y = torch.randn(200, 1)                  # TODO: 标签；分类任务改成 int64 的类别下标
    # 特征必须是 float32，分类任务的标签必须是 int64，这两条类型错了直接报 RuntimeError
    return TensorDataset(x.float(), y.float())


# ---------------------------------------------------------------- 模型
class MyModel(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # TODO: 按任务堆层；层与层之间的维度要接得上
        self.linear1 = nn.Linear(in_features, 64)
        self.out = nn.Linear(64, out_features)

    def forward(self, x):
        x = torch.relu(self.linear1(x))
        # TODO: 多分类时这里不要加 softmax，交给 CrossEntropyLoss
        return self.out(x)


# ---------------------------------------------------------------- 训练
def train(dataset, in_features, out_features,
          epochs=50, batch_size=16, lr=1e-3, seed=0):
    torch.manual_seed(seed)                  # 固定种子，结果才可复现
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = MyModel(in_features, out_features)
    criterion = nn.MSELoss()                 # TODO: 分类换 nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    for epoch in range(epochs):
        model.train()
        total_loss, sample_num = 0.0, 0
        for bx, by in dataloader:
            y_pred = model(bx)
            loss = criterion(y_pred, by)

            optimizer.zero_grad()            # 顺序不能乱：清零 -> 反向 -> 更新
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * len(by)
            sample_num += len(by)
        print('epoch %3d  loss %.5f' % (epoch + 1, total_loss / sample_num))

    torch.save(model.state_dict(), 'model.pt')   # 存权重而不是整个模型对象
    return model


if __name__ == '__main__':
    ds = create_dataset()
    train(ds, in_features=4, out_features=1)

"""形状排错速查：把最常见的五种维度报错原地复现一遍，并给出修法。

跑一次这份文件，以后再看到同样的报错就能直接定位。
每一段都用 try/except 抓住报错并打印，不会中断。
"""
import torch
from torch import nn


def case(title):
    print('\n' + '=' * 56)
    print(title)
    print('=' * 56)


# ---------------------------------------------------------------- 一、少了 batch 维
case('一、单条样本直接喂进网络')
layer = nn.Linear(4, 2)
one_sample = torch.randn(4)            # (4,) 少了 batch 维
try:
    print(layer(one_sample).shape)     # 一维输入其实能跑，但语义上是 batch=1 被隐式处理
except RuntimeError as e:
    print('报错:', e)
print('规范写法:', layer(one_sample.unsqueeze(0)).shape)   # (1, 2)


# ---------------------------------------------------------------- 二、特征数对不上
case('二、输入特征数与层定义不一致')
wrong_feat = torch.randn(8, 5)         # 层要 4 个特征，给了 5 个
try:
    layer(wrong_feat)
except RuntimeError as e:
    print('报错:', e)
print('修法：要么改层的 in_features，要么检查数据列数是不是多带了 id/标签列')


# ---------------------------------------------------------------- 三、矩阵乘法接不上
case('三、矩阵乘法维度接不上')
m1 = torch.randn(3, 2)
m2 = torch.randn(3, 2)
try:
    m1 @ m2
except RuntimeError as e:
    print('报错:', e)
print('修法:', (m1 @ m2.T).shape, '—— (3,2)@(2,3) 才接得上')


# ---------------------------------------------------------------- 四、view 遇到非连续内存
case('四、transpose 之后直接 view')
t = torch.arange(6).reshape(2, 3)
swapped = t.transpose(0, 1)
print('是否内存连续:', swapped.is_contiguous())
try:
    swapped.view(2, 3)
except RuntimeError as e:
    print('报错:', e)
print('修法一 contiguous:', swapped.contiguous().view(2, 3).shape)
print('修法二 直接 reshape:', swapped.reshape(2, 3).shape)


# ---------------------------------------------------------------- 五、不报错但算错：广播
case('五、最危险的一种：不报错，但算错')
y_pred = torch.randn(8, 1)
y_true = torch.randn(8)
loss_wrong = nn.MSELoss()(y_pred, y_true)
loss_right = nn.MSELoss()(y_pred, y_true.reshape(-1, 1))
print('两边形状:', y_pred.shape, y_true.shape)
print('广播后的差值形状:', (y_pred - y_true).shape, '<- 本该是 (8, 1)')
print('错的 loss %.6f    对的 loss %.6f' % (loss_wrong.item(), loss_right.item()))
print('修法：算损失前 assert 两边形状相同')