如何在PyTorch实现的CNN食物分类任务中添加数据归一化？

我是CNN初学者，正在基于食物分类任务学习CNN，以下是我的初始实现代码。在DATASET部分，我将训练集和验证集从numpy数组转换为tensor，此时tensor的形状为[9866, 128, 128, 3]。由于3通道需要放在第一维度，我使用transpose方法调整了维度顺序，之后用Data.TensorDataset拼接训练数据和训练标签，使用Data.DataLoader是为了按批次加载以提升训练速度。

import os
import numpy as np
import pandas as pd
import cv2
import torch
import torch.nn as nn
from torch.nn import functional as F
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch import optim
import pandas as pd
from torch.utils.data import DataLoader, Dataset
import torch.utils.data as Data



'''Initialize Params'''
epochs = 3
learning_rate = 0.0001
momentum = 0.5
batch_size = 128


'''Load Data'''
def readFile(path,label):
    image_dir = sorted(os.listdir(path))
    # x stores photos
    x = np.zeros((len(image_dir),128,128,3),dtype=np.uint8)
    # y stores labels
    y = np.zeros((len(image_dir)), dtype=np.uint8)
    for i, file in enumerate(image_dir):
        img = cv2.imread(os.path.join(path, file))
        x[i, :, :] = cv2.resize(img,(128, 128))
        if label:
            y[i] = int(file.split("_")[0])
    if label:
        return x,y
    else:
        return x

train_x, train_y = readFile('./food/training',True)
val_x, val_y = readFile('./food/validation',True)
test_x = readFile('./food/testing',False)
# print("Reading data: ")
# print("Size of training data = {}".format(len(train_x)))
# print("Size of validation data = {}".format(len(val_x)))
# print("Size of Testing data = {}".format(len(test_x)))


'''DataSet'''
train_x = torch.tensor(train_x)
# print(train_x.shape)
train_x = train_x.transpose(1,3).float()
train_y = torch.tensor(train_y)
val_x = torch.tensor(val_x)
val_x = val_x.transpose(1, 3).float()
val_y = torch.tensor(val_y)


train_dataset = Data.TensorDataset(train_x,train_y)
val_dataset = Data.TensorDataset(val_x,val_y)

train_loader = Data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)
val_loader = Data.DataLoader(dataset=val_dataset,batch_size=batch_size,shuffle=True) 

目前我的模型训练集准确率仅为68%，希望进一步提升效果。我在网上搜索得知可以添加数据归一化操作，但只找到了如下实现方式：

transform = transforms.Compose([
    transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
    ]
)

我不清楚如何将其和Data.DataLoader结合使用。我还查到了另一种将numpy数据转换为dataloader的自定义Dataset实现，实现代码如下：

train_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomHorizontalFlip(), 
    transforms.RandomRotation(15), 
    transforms.ToTensor(), 
])

test_transform = transforms.Compose([
    transforms.ToPILImage(),                                    
    transforms.ToTensor(),
])


class ImgDataset(Dataset):
    def __init__(self, x, y=None, transform=None):
        self.x = x
        self.y = y
        if y is not None:
            self.y = torch.LongTensor(y)
        self.transform = transform
    def __len__(self):
        return len(self.x)
    def __getitem__(self, index):
        X = self.x[index]
        if self.transform is not None:
            X = self.transform(X)
        if self.y is not None:
            Y = self.y[index]
            return X, Y
        else:
            return X

train_set = ImgDataset(train_x, train_y, train_transform)
val_set = ImgDataset(val_x, val_y, test_transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False)

我的实现方式可能不够灵活，但我希望基于现有方案调整，恳请各位提供帮助，感谢解答。完整实现代码如下：

import os
import numpy as np
import pandas as pd
import cv2
import torch
import torch.nn as nn
from torch.nn import functional as F
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch import optim
import pandas as pd
from torch.utils.data import DataLoader, Dataset
import torch.utils.data as Data



'''Initialize Params'''
epochs = 3
learning_rate = 0.0001
momentum = 0.5
batch_size = 128


transform = transforms.Compose([
    transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
    ]
)



'''Load Data'''
def readFile(path,label):
    image_dir = sorted(os.listdir(path))
    # x stores photos
    x = np.zeros((len(image_dir),128,128,3),dtype=np.uint8)
    # y stores labels
    y = np.zeros((len(image_dir)), dtype=np.uint8)
    for i, file in enumerate(image_dir):
        img = cv2.imread(os.path.join(path, file))
        x[i, :, :] = cv2.resize(img,(128, 128))
        if label:
            y[i] = int(file.split("_")[0])
    if label:
        return x,y
    else:
        return x

train_x, train_y = readFile('./food/training',True)
val_x, val_y = readFile('./food/validation',True)
test_x = readFile('./food/testing',False)
# print("Reading data: ")
# print("Size of training data = {}".format(len(train_x)))
# print("Size of validation data = {}".format(len(val_x)))
# print("Size of Testing data = {}".format(len(test_x)))


'''DataSet'''
train_x = torch.tensor(train_x)
# print(train_x.shape)
train_x = train_x.transpose(1,3).float()
train_y = torch.tensor(train_y)
val_x = torch.tensor(val_x)
val_x = val_x.transpose(1, 3).float()
val_y = torch.tensor(val_y)


train_dataset = Data.TensorDataset(train_x,train_y)
val_dataset = Data.TensorDataset(val_x,val_y)

train_loader = Data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)
val_loader = Data.DataLoader(dataset=val_dataset,batch_size=batch_size,shuffle=True) 


'''Create Model'''
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # nn.Conv2d(input_channel, output_channel, kernel, stride)
        self.conv1 = nn.Conv2d(3,64,5,1,1) 
        nn.BatchNorm2d(64) 
        self.conv2 = nn.Conv2d(64,128,5,1,1)
        nn.BatchNorm2d(128)
        self.conv3 = nn.Conv2d(128,256,5,1,1)
        nn.BatchNorm2d(256)
        self.conv4 = nn.Conv2d(256,256,5,1,1)
        nn.BatchNorm2d(256)
        self.conv4_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(6*6*256, 1024)  
        # self.fc1 = nn.Linear(512*4*4, 1024)
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.fc4 = nn.Linear(256, 11)


    def forward(self, x):
        # maxpooling 1
        x = self.conv1(x)
        x = F.relu(x) # 124*124*64
        x = F.max_pool2d(x, 2) # 62*62*20

        # maxpooling 2
        x = self.conv2(x)
        x = F.relu(x) # 58*58*128
        x = F.max_pool2d(x, 2) # 29*29*40

        # maxpooling 3
        x = self.conv3(x)
        x = F.relu(x) # 25*25*256
        x = F.max_pool2d(x, 2) # 12*12*100

        # maxpooling 4
        x = self.conv4(x)
        x = F.relu(x) # 8*8*256
        x = F.max_pool2d(x, 2) # 4*4*256

        x = x.view(-1,6*6*256)
        # print(x.shape)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        # return F.log_softmax(x)
        return F.log_softmax(x)


'''Initialize the network'''
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=learning_rate)


'''Train Model'''
if __name__ == '__main__':
    # print(train_x.shape)  torch.Size([9866, 128, 128, 3])
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    net.to(device)
    for epoch in range(epochs):
        for i,data in enumerate(train_loader):
            x,y = data
            # print(x)
            # print('--------------------')
            sum_loss = 0.0
            optimizer.zero_grad()  
            x,y = Variable(x).cuda(), Variable(y).cuda()
            outputs = net(x)  
            loss = criterion(outputs, y.long())  
            loss.backward()
            optimizer.step()  
            # print(loss.item())
            # sum_loss += loss.item()
            # if i % 500 == 99:
            #     print('[%d,%d] loss:%.03f' %
            #             (epoch + 1, i + 1, sum_loss / 500))
            #     sum_loss = 0.0


        '''Test Model with Validation'''
        net.eval() 
        correct = 0
        total = 0
        for val_data in val_loader:
            x,y = val_data
            x, y = Variable(x).cuda(), Variable(y).cuda()
            output_test = net(x)
            _, predicted = torch.max(output_test, 1)
            total += y.size(0)
            correct += (predicted == y).sum()
        print("correct1: ", correct)
        print("Test acc: {0}".format(correct.item() /len(val_loader)))

可直接落地的调整方案 #

1. 数据加载模块调整（支持归一化+数据增强，兼容你现有逻辑） #

删除原来DataSet模块下手动转tensor、调整维度的代码，直接复用你找到的自定义Dataset类即可，transform会自动完成维度调整、[0,255]到[0,1]的归一化，还可以额外加数据增强提升泛化能力：

# 替换原有transform和DataSet模块代码
train_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.RandomHorizontalFlip(), # 随机水平翻转
    transforms.RandomRotation(15), # 随机旋转15度以内
    transforms.ToTensor(), # 自动完成归一化+维度调整
    # 可选：追加ImageNet预训练分布归一化，适合后续微调预训练模型
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.ToPILImage(),                                    
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

class ImgDataset(Dataset):
    def __init__(self, x, y=None, transform=None):
        self.x = x
        self.y = y
        if y is not None:
            self.y = torch.LongTensor(y)
        self.transform = transform
    def __len__(self):
        return len(self.x)
    def __getitem__(self, index):
        X = self.x[index]
        if self.transform is not None:
            X = self.transform(X)
        if self.y is not None:
            Y = self.y[index]
            return X, Y
        else:
            return X

# 直接用readFile返回的numpy数组创建数据集即可
train_set = ImgDataset(train_x, train_y, train_transform)
val_set = ImgDataset(val_x, val_y, test_transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False)

2. 现有模型和训练逻辑的bug修复（直接影响准确率） #

• 你定义的nn.BatchNorm2d没有赋值给类属性，相当于没有生效，需要改成self.bn1 = nn.BatchNorm2d(64)的形式并在forward中调用
• 全连接层输入维度计算错误：四次池化后输出维度是44256，不是66256，维度不匹配会导致训练异常
• 验证准确率计算错误：原有代码是除以loader的批次数量，应该除以验证集总样本数total
• 训练轮次设置过小，3轮不足以让模型收敛，建议调到20-30轮，可以配合学习率衰减进一步提升效果
• F.log_softmax需要指定dim参数，避免计算异常：F.log_softmax(x, dim=1)

修正后的模型代码如下：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3,64,5,1,1) 
        self.bn1 = nn.BatchNorm2d(64) 
        self.conv2 = nn.Conv2d(64,128,5,1,1)
        self.bn2 = nn.BatchNorm2d(128)
        self.conv3 = nn.Conv2d(128,256,5,1,1)
        self.bn3 = nn.BatchNorm2d(256)
        self.conv4 = nn.Conv2d(256,256,5,1,1)
        self.bn4 = nn.BatchNorm2d(256)
        self.conv4_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(4*4*256, 1024)  
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.fc4 = nn.Linear(256, 11)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)

        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)

        x = self.conv3(x)
        x = self.bn3(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)

        x = self.conv4(x)
        x = self.bn4(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)

        x = x.view(-1,4*4*256)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return F.log