{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验 01：张量、梯度与优化\n",
    "\n",
    "**运行条件：** CPU，约 1 分钟，内存低于 1 GB。\n",
    "\n",
    "本实验不调用高级训练框架。目标是观察张量形状、计算图、链式法则和参数更新之间的关系。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(7)\n",
    "print('PyTorch:', torch.__version__)\n",
    "print('device: cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 从形状开始\n",
    "\n",
    "LLM 中最常见的激活形状是 `[batch, sequence, hidden]`。线性层只改变最后一个维度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch, sequence, hidden, output = 2, 4, 8, 16\n",
    "x = torch.randn(batch, sequence, hidden)\n",
    "w = torch.randn(hidden, output)\n",
    "y = x @ w\n",
    "print('x:', tuple(x.shape))\n",
    "print('w:', tuple(w.shape))\n",
    "print('y:', tuple(y.shape))\n",
    "assert y.shape == (batch, sequence, output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 验证链式法则\n",
    "\n",
    "对标量函数 $y=(x^2+1)^3$，解析梯度为 $6x(x^2+1)^2$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor(2.0, requires_grad=True)\n",
    "y = (x**2 + 1)**3\n",
    "y.backward()\n",
    "expected = 6 * x.detach() * (x.detach()**2 + 1)**2\n",
    "print('autograd:', x.grad.item())\n",
    "print('analytic:', expected.item())\n",
    "assert torch.allclose(x.grad, expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 用梯度下降拟合直线\n",
    "\n",
    "合成数据满足 $y=3x-0.7+噪声$。实验记录每一步损失，并检查参数是否收敛到真实值附近。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = torch.linspace(-2, 2, 128).unsqueeze(1)\n",
    "targets = 3.0 * features - 0.7 + 0.08 * torch.randn_like(features)\n",
    "weight = torch.randn(1, 1, requires_grad=True)\n",
    "bias = torch.zeros(1, requires_grad=True)\n",
    "losses = []\n",
    "for step in range(160):\n",
    "    prediction = features @ weight + bias\n",
    "    loss = ((prediction - targets) ** 2).mean()\n",
    "    loss.backward()\n",
    "    with torch.no_grad():\n",
    "        weight -= 0.08 * weight.grad\n",
    "        bias -= 0.08 * bias.grad\n",
    "        weight.grad.zero_()\n",
    "        bias.grad.zero_()\n",
    "    losses.append(loss.item())\n",
    "print(f'weight={weight.item():.3f}, bias={bias.item():.3f}, loss={losses[-1]:.5f}')\n",
    "plt.plot(losses)\n",
    "plt.yscale('log')\n",
    "plt.xlabel('step'); plt.ylabel('MSE'); plt.title('Optimization trajectory'); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 练习\n",
    "\n",
    "1. 把学习率改成 `0.8`，解释损失为何震荡或发散。\n",
    "2. 删除 `zero_()`，观察梯度累积造成的影响。\n",
    "3. 将输入扩展到两个特征，写出对应的权重形状。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}