神经网络对于的调参主要指超参数的调参，在众多超参数之中，批量大小（Batch Size）占有着举足轻重的作用。理论上，批量大小决定着单次送进神经网络中的样本规模，合理的批量大小可以充分的利用 GPU 的并行计算能力。本文主要通过 Fashion-Mnist 数据集，探求不同批量大小对于收敛的影响。

测试代码

首先编写一个基准代码，用于测试不同的batch size 对于收敛的影响：

import numpy as np
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torch import nn
import time

device = torch.device('cuda:0')

trans = transforms.ToTensor()

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

class Mnist_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1)

    def forward(self, xb):
        xb = xb.view(-1, 1, 28, 28)
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4)
        return xb.view(-1, xb.size(1))

def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)

def fit(epochs, model, loss_func, opt, train_dl, valid_dl, results, batch_size):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            xb = xb.to(device)
            yb = yb.to(device)
            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb.to(device), yb.to(device)) for xb, yb in valid_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)
        results[batch_size].append(val_loss)

以上是通过 CNN 对 Fashion-Mnist 进行分类的简单代码，results 用于保存中间结果。

测试不同的批量大小

batch_size_range = [16,32,64,128,256] # 
epochs = 20
results = {}

for batch_size in batch_size_range:
    start = time.perf_counter()
    model = Mnist_CNN().to(device)
    train_dl = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=8, prefetch_factor=32)
    valid_dl = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=8, prefetch_factor=32)

    loss = nn.CrossEntropyLoss()
    optim = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

    results[batch_size] = []
    fit(epochs, model, loss, optim, train_dl, valid_dl, results, batch_size)
    results[batch_size].append(time.perf_counter()-start)

这里我们测试了 16,32,64,128,256 的批量大小，模型在更换batch size训练时重新初始化，避免前后影响。

使用 matplotlib进行可视化：

import matplotlib.pyplot as plt

x = range(1, epochs + 1)

fig, ax = plt.subplots()
for item in results:
    ax.plot(list(x), results[item][:-1], label=str(item))

ax.set_xlabel('epochs') 
ax.set_ylabel('val loss') 
ax.set_title('different Batch Size')
ax.legend()

plt.show()

代码比较简单，不予赘述，直接看结果：

不同批量大小下收敛情况

40 epoch 结果如下：

40 epochs 下不同批量大小下收敛情况

收敛时间：

for item in results:
    print(f"batch size = {item}, 20 epochs speend {results[item][-1]}")

不同批量大小的收敛时间

结论

综上所述，BS（Batch Size）越大，在GPU上训练时越能发挥并行计算的能力，也就是硬件利用率越高，但模型收敛所需的 epochs 也随之越多，反之，硬件利用率不高，收敛时间长，但是模型可以在同一个epochs中进行更多的参数校准，模型收敛也会更快。在实际调参过程中，BS的取值既不能太小也不能太大，选择折中的数值都是可行的，并没有一个最优值。不过，BS 不会影响模型的拟合能力。