# 先繼續之前的代碼
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader , Dataset # DataLoader 是 PyTorch 中用于加載數據的工具
from torchvision import datasets, transforms # torchvision 是一個用于計算機視覺的庫，datasets 和 transforms 是其中的模塊
import matplotlib.pyplot as plt
# 設置隨機種子，確保結果可復現
torch.manual_seed(42)# 1. 數據預處理，該寫法非常類似于管道pipeline
# transforms 模塊提供了一系列常用的圖像預處理操作# 先歸一化，再標準化
transform = transforms.Compose([transforms.ToTensor(),  # 轉換為張量并歸一化到[0,1]transforms.Normalize((0.1307,), (0.3081,))  # MNIST數據集的均值和標準差，這個值很出名，所以直接使用
])
import matplotlib.pyplot as plt# 2. 加載MNIST數據集，如果沒有會自動下載
train_dataset = datasets.MNIST(root='./data',train=True,download=True,transform=transform
)test_dataset = datasets.MNIST(root='./data',train=False,transform=transform
)# 隨機選擇一張圖片，可以重復運行，每次都會隨機選擇
sample_idx = torch.randint(0, len(train_dataset), size=(1,)).item() # 隨機選擇一張圖片的索引
# len(train_dataset) 表示訓練集的圖片數量；size=(1,)表示返回一個索引；torch.randint() 函數用于生成一個指定范圍內的隨機數,item() 方法將張量轉換為 Python 數字
image, label = train_dataset[sample_idx] # 獲取圖片和標簽
# 可視化原始圖像（需要反歸一化）
def imshow(img):img = img * 0.3081 + 0.1307  # 反標準化npimg = img.numpy()plt.imshow(npimg[0], cmap='gray') # 顯示灰度圖像plt.show()print(f"Label: {label}")
imshow(image)# 打印下圖片的形狀
image.shape#torch.Size([1, 28, 28])

# 打印一張彩色圖像，用cifar-10數據集
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np# 設置隨機種子確保結果可復現
torch.manual_seed(42)
# 定義數據預處理步驟
transform = transforms.Compose([transforms.ToTensor(),  # 轉換為張量并歸一化到[0,1]transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 標準化處理
])# 加載CIFAR-10訓練集
trainset = torchvision.datasets.CIFAR10(root='./data',train=True,download=True,transform=transform
)# 創建數據加載器
trainloader = torch.utils.data.DataLoader(trainset,batch_size=4,shuffle=True
)# CIFAR-10的10個類別
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')# 隨機選擇一張圖片
sample_idx = torch.randint(0, len(trainset), size=(1,)).item()
image, label = trainset[sample_idx]# 打印圖片形狀
print(f"圖像形狀: {image.shape}")  # 輸出: torch.Size([3, 32, 32])
print(f"圖像類別: {classes[label]}")# 定義圖像顯示函數（適用于CIFAR-10彩色圖像）
def imshow(img):img = img / 2 + 0.5  # 反標準化處理，將圖像范圍從[-1,1]轉回[0,1]npimg = img.numpy()plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 調整維度順序：(通道,高,寬) → (高,寬,通道)plt.axis('off')  # 關閉坐標軸顯示plt.show()# 顯示圖像
imshow(image)

# 先歸一化，再標準化
transform = transforms.Compose([transforms.ToTensor(),  # 轉換為張量并歸一化到[0,1]transforms.Normalize((0.1307,), (0.3081,))  # MNIST數據集的均值和標準差，這個值很出名，所以直接使用
])
import matplotlib.pyplot as plt# 2. 加載MNIST數據集，如果沒有會自動下載
train_dataset = datasets.MNIST(root='./data',train=True,download=True,transform=transform
)test_dataset = datasets.MNIST(root='./data',train=False,transform=transform
)# 定義兩層MLP神經網絡
class MLP(nn.Module):def __init__(self):super(MLP, self).__init__()self.flatten = nn.Flatten()  # 將28x28的圖像展平為784維向量self.layer1 = nn.Linear(784, 128)  # 第一層：784個輸入，128個神經元self.relu = nn.ReLU()  # 激活函數self.layer2 = nn.Linear(128, 10)  # 第二層：128個輸入，10個輸出（對應10個數字類別）def forward(self, x):x = self.flatten(x)  # 展平圖像x = self.layer1(x)   # 第一層線性變換x = self.relu(x)     # 應用ReLU激活函數x = self.layer2(x)   # 第二層線性變換，輸出logitsreturn x# 初始化模型
model = MLP()device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)  # 將模型移至GPU（如果可用）from torchsummary import summary  # 導入torchsummary庫
print("\n模型結構信息：")
summary(model, input_size=(1, 28, 28))  # 輸入尺寸為MNIST圖像尺寸

我們關注和之前結構化MLP的差異

1. 輸入需要展平操作

MLP 的輸入層要求輸入是一維向量，但 MNIST 圖像是二維結構（28×28 像素），形狀為 [1, 28, 28]（通道 × 高 × 寬）。nn.Flatten()展平操作 將二維圖像 “拉成” 一維向量（784=28×28 個元素），使其符合全連接層的輸入格式。

其中不定義這個flatten方法，直接在前向傳播的過程中用 x = x.view(-1, 28 * 28) 將圖像展平為一維向量也可以實現

1. 輸入數據的尺寸包含了通道數input_size=(1, 28, 28)

2. 參數的計算

• 第一層 layer1（全連接層）

權重參數：輸入維度 × 輸出維度 = 784 × 128 = 100,352

偏置參數：輸出維度 = 128

合計：100,352 + 128 = 100,480

• 第二層 layer2（全連接層）

權重參數：輸入維度 × 輸出維度 = 128 × 10 = 1,280

偏置參數：輸出維度 = 10
合計：1,280 + 10 = 1,290

• 總參數：100,480（layer1） + 1,290（layer2） = 101,770

2.2 彩色圖像模型的定義

class MLP(nn.Module):def __init__(self, input_size=3072, hidden_size=128, num_classes=10):super(MLP, self).__init__()# 展平層：將3×32×32的彩色圖像轉為一維向量# 輸入尺寸計算：3通道 × 32高 × 32寬 = 3072self.flatten = nn.Flatten()# 全連接層self.fc1 = nn.Linear(input_size, hidden_size)  # 第一層self.relu = nn.ReLU()self.fc2 = nn.Linear(hidden_size, num_classes)  # 輸出層def forward(self, x):x = self.flatten(x)  # 展平：[batch, 3, 32, 32] → [batch, 3072]x = self.fc1(x)      # 線性變換：[batch, 3072] → [batch, 128]x = self.relu(x)     # 激活函數x = self.fc2(x)      # 輸出層：[batch, 128] → [batch, 10]return x# 初始化模型
model = MLP()device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)  # 將模型移至GPU（如果可用）from torchsummary import summary  # 導入torchsummary庫
print("\n模型結構信息：")
summary(model, input_size=(3, 32, 32))  # CIFAR-10 彩色圖像（3×32×32）

模型結構信息：
----------------------------------------------------------------
? ? ? ? Layer (type) ? ? ? ? ? ? ? Output Shape ? ? ? ? Param #
================================================================
? ? ? ? ? ?Flatten-1 ? ? ? ? ? ? ? ? [-1, 3072] ? ? ? ? ? ? ? 0
? ? ? ? ? ? Linear-2 ? ? ? ? ? ? ? ? ?[-1, 128] ? ? ? ? 393,344
? ? ? ? ? ? ? ReLU-3 ? ? ? ? ? ? ? ? ?[-1, 128] ? ? ? ? ? ? ? 0
? ? ? ? ? ? Linear-4 ? ? ? ? ? ? ? ? ? [-1, 10] ? ? ? ? ? 1,290
================================================================
Total params: 394,634
Trainable params: 394,634
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.01
Forward/backward pass size (MB): 0.03
Params size (MB): 1.51
Estimated Total Size (MB): 1.54
----------------------------------------------------------------

• 第一層 layer1（全連接層）

權重參數：輸入維度 × 輸出維度 = 3072 × 128 = 393,216

偏置參數：輸出維度 = 128

合計：393,216 + 128 = 393,344

• 第二層 layer2（全連接層）

權重參數：輸入維度 × 輸出維度 = 128 × 10 = 1,280

偏置參數：輸出維度 = 10

合計：1,280 + 10 = 1,290

• 總參數：393,344（layer1） + 1,290（layer2） = 394,634

2.3 模型定義與batchsize的關系

實際定義中，輸入圖像還存在batchsize這一維度

在 PyTorch 中，模型定義和輸入尺寸的指定不依賴于 batch_size，無論設置多大的 batch_size，模型結構和輸入尺寸的寫法都是不變的。

class MLP(nn.Module):def __init__(self):super().__init__()self.flatten = nn.Flatten() # nn.Flatten()會將每個樣本的圖像展平為 784 維向量，但保留 batch 維度。self.layer1 = nn.Linear(784, 128)self.relu = nn.ReLU()self.layer2 = nn.Linear(128, 10)def forward(self, x):x = self.flatten(x)  # 輸入：[batch_size, 1, 28, 28] → [batch_size, 784]x = self.layer1(x)   # [batch_size, 784] → [batch_size, 128]x = self.relu(x)x = self.layer2(x)   # [batch_size, 128] → [batch_size, 10]return x

PyTorch 模型會自動處理 batch 維度（即第一維），無論 batch_size 是多少，模型的計算邏輯都不變。batch_size 是在數據加載階段定義的，與模型結構無關。

summary(model, input_size=(1, 28, 28))中的input_size不包含 batch 維度，只需指定樣本的形狀（通道 × 高 × 寬）。

總結：batch_size與模型定義的關系**

三、顯存占用的主要組成部分

昨天說到了在面對數據集過大的情況下，由于無法一次性將數據全部加入到顯存中，所以采取了分批次加載這種方式。即一次只加載一部分數據，保證在顯存的范圍內。

那么顯存設置多少合適呢？如果設置的太小，那么每個batchsize的訓練不足以發揮顯卡的能力，浪費計算資源；如果設置的太大，會出現OOT（out of memory）

顯存一般被以下內容占用：

1. 模型參數與梯度：模型的權重（Parameters）和對應的梯度（Gradients）會占用顯存，尤其是深度神經網絡（如 Transformer、ResNet 等），一個 1 億參數的模型（如 BERT-base），單精度（float32）參數占用約 400MB（1e8×4Byte），加上梯度則翻倍至 800MB（每個權重參數都有其對應的梯度）。

2. 部分優化器（如 Adam）會為每個參數存儲動量（Momentum）和平方梯度（Square Gradient），進一步增加顯存占用（通常為參數大小的 2-3 倍）

3. 其他開銷。

下面以手寫數據集為例

from torch.utils.data import DataLoader# 定義訓練集的數據加載器，并指定batch_size
train_loader = DataLoader(dataset=train_dataset,  # 加載的數據集batch_size=64,          # 每次加載64張圖像shuffle=True            # 訓練時打亂數據順序
)# 定義測試集的數據加載器（通常batch_size更大，減少測試時間）
test_loader = DataLoader(dataset=test_dataset,batch_size=1000,shuffle=False
)

手寫數據集（MNIST）和當前 MLP 模型，顯存占用的計算可以簡化為以下幾個部分。

3.1 模型參數與梯度（FP32 精度）

參數總量：101,770 個參數

• 1字節（Byte）= 8位（bit），是計算機存儲的最小尋址單位。
• 位（bit）是二進制數的最小單位（0或1），例如0b1010表示4位二進制數。
• 1KB=1024字節；1MB=1024KB=1,048,576字節

常見數據類型的字節占用

• MNIST數據集的原始圖像像素值為0-255的整數（uint8類型，占1字節），表示灰度值（0=黑，255=白）。
• 但PyTorch的transforms.ToTensor()會將其歸一化到[0, 1]范圍，并轉換為?float32類型（浮點型更適合神經網絡計算）。

計算示例：單張MNIST圖像的顯存占用

1. 原始像素值（uint8，未轉換時）

   • 尺寸：28×28像素
   • 單像素占用：1字節（uint8）
   • 總占用：28×28×1 = 784字節 ≈ 0.766 KB

2. 轉換為float32張量后

   • 尺寸：1×28×28（通道×高×寬）
   • 單像素占用：4字節（float32）
   • 總占用：1×28×28×4 = 3136字節 ≈ 3.06 KB

單精度（float32）參數占用：101,770 × 4 Byte ≈ 403 KB

梯度是損失函數對模型參數的導數（?Loss/?Weight），用于指示參數更新的方向和幅度。梯度是損失函數對模型參數的導數（?Loss/?Weight），用于指示參數更新的方向和幅度。因此在默認情況下，梯度的數據類型和數目與參數相同。

梯度占用（反向傳播時）：與參數相同，合計約 806 KB

3.2?優化器狀態

SGD

• SGD優化器不存儲額外動量，因此無額外顯存占用。
• SGD 隨機梯度下降，最基礎的優化器，直接沿梯度反方向更新參數。
• 參數更新公式：w = w - learning_rate * gradient

Adam

• Adam優化器：自適應學習率優化器，結合了動量（Momentum）和梯度平方的指數移動平均。
• 每個參數存儲動量（m）和平方梯度（v），占用約?101,770 × 8 Byte ≈ 806 KB
• 動量（m）：每個參數對應一個動量值，數據類型與參數相同（float32），占用 403 KB。
• 梯度平方（v）：每個參數對應一個梯度平方值，數據類型與參數相同（float32），占用 403 KB。

3.3?數據批量（batch_size）的顯存占用

• 單張圖像尺寸：1×28×28（通道×高×寬），歸一化轉換為張量后為float32類型
  • 單張圖像顯存占用：1×28×28×4 Byte = 3,136 Byte ≈ 3 KB
• 批量數據占用：batch_size × 單張圖像占用
  • 例如：batch_size=64?時，數據占用為?64×3 KB ≈ 192 KB
  • batch_size=1024?時，數據占用為?1024×3 KB ≈ 3 MB

3.4.?前向/反向傳播中間變量

• 對于兩層MLP，中間變量（如layer1的輸出）占用較小：
  • batch_size×128維向量：batch_size×128×4 Byte = batch_size×512 Byte
  • 例如batch_size=1024時，中間變量約?512 KB

以SGD為例，此時其他參數占用固定，batchsize會影響顯存占用

在 PyTorch 中，在使用DataLoader加載數據時，如果不指定batch_size參數，默認值是1

，即每次迭代返回一個樣本。這與一次性使用全部數據進行訓練是完全不同的概念。如果想要一次性使用全部數據進行訓練，需要手動將batch_size設置為數據集的大小，但對于大型數據集，這樣做通常會導致內存不足，因為一次性將所有數據加載到內存中可能會超出硬件的內存限制。

大規模數據時，通常從16開始測試，然后逐漸增加，確保代碼運行正常且不報錯，直到出現 內存不足（OOM）報錯 或訓練效果下降，此時選擇略小于該值的 batch_size。

訓練時候搭配 nvidia-smi 監控顯存占用，合適的 batch_size = 硬件顯存允許的最大值 × 0.8（預留安全空間），并通過訓練效果驗證調整。

補充說明： batchsize對于訓練的影響

在深度學習中，使用較大的 batch_size（批量大小）相比單樣本訓練（batch_size=1）有以下核心優勢

• 并行計算能力最大化，減小訓練時間；且大幅減少更新次數
• 梯度方向更準確，單樣本訓練的梯度僅基于單個數據點，可能包含大量噪聲（尤其是數據分布不均或存在異常值時）。大 batch_size 的梯度是多個樣本的平均值，能抵消單個樣本的隨機性，梯度方向更接近真實分布的 “全局最優方向”。會讓訓練過程更穩定，波動更小

@浙大疏錦行