import torch from net import simpleconv from torchvision import transforms, datasets

损失函数(用于更新参数)、优化器(梯度下降以一个什么样的方式下降)、学习率

import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler from tensorboardX import SummaryWriter import time import os

def train(train_loader, model, criterion, optimizer, device, len_train, batch_size): num_loss = 0.0 # 累计损失值 num_corrects = 0.0 # 累计正确分类的样本数 model.train() # 设置模型为训练模式

for i, (data, target) in enumerate(train_loader):  # 遍历训练数据加载器
    data = data.to(device)  # 将输入数据移动到指定设备(如GPU)
    target = target.to(device)  # 将目标标签移动到指定设备

    optimizer.zero_grad()  # 梯度清零,避免累积梯度
    outputs = model(data)  # 前向传播,得到模型的输出
    _, preds = torch.max(outputs, 1)  # 获取预测结果中的最大值及其对应的索引
    loss = criterion(outputs, target)  # 计算损失值
    loss.backward()  # 反向传播,计算参数的梯度
    optimizer.step()  # 根据梯度更新参数

    num_corrects += torch.sum(preds == target).item()  # 统计准确分类的样本数
    num_loss = num_loss + loss.item()  # 累计损失值

train_loss = num_loss / (len_train // batch_size + 1)  # 计算每个epoch的平均训练损失
train_acc = num_corrects / len_train  # 计算每个epoch的训练准确率

return train_loss, train_acc  # 返回每个epoch的平均训练损失和训练准确率

def val(val_loader, model, criterion, device, len_val, batch_size): num_loss = 0.0 # 累计损失值 num_corrects = 0.0 # 累计正确分类的样本数 model.eval() # 将模型切换到验证模式

with torch.no_grad():  # 在验证阶段,不进行梯度计算和参数更新
    for i, (data, target) in enumerate(val_loader):  # 遍历验证数据加载器
        data = data.to(device)  # 将输入数据移动到指定设备(如GPU)
        target = target.to(device)  # 将目标标签移动到指定设备

        outputs = model(data)  # 前向传播,得到模型的输出
        _, preds = torch.max(outputs.data, 1)  # 获取预测结果中的最大值及其对应的索引
        loss = criterion(outputs, target)  # 计算损失值

        num_corrects += torch.sum(preds == target).item()  # 统计准确分类的样本数
        num_loss = num_loss + loss.item()  # 累计损失值

val_loss = num_loss / (len_val // batch_size + 1)  # 计算每个epoch的平均验证损失
val_acc = num_corrects / len_val  # 计算每个epoch的验证准确率

return val_loss, val_acc  # 返回每个epoch的平均验证损失和验证准确率

class simpleconv(nn.Module): def init(self, nclass, dropout_rate=0.5): super(simpleconv, self).init() self.conv1 = nn.Sequential( nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.conv2 = nn.Sequential( nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.fc = nn.Linear(32 * 12 * 12, nclass) self.dropout = nn.Dropout(dropout_rate) # 添加dropout层

def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)
    x = self.dropout(x)  # 添加dropout层
    return x

batch_size = 100 nclass = 13 num_epochs = 364 data_dir = './testimage'

if torch.cuda.is_available(): device = "cuda" else: device = "cpu" model = simpleconv(nclass, dropout_rate=0.5).to(device)

train_transforms = transforms.Compose([ transforms.Resize((64, 64)), transforms.RandomResizedCrop(48), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

val_transforms = transforms.Compose([ transforms.Resize((64, 64)), transforms.CenterCrop(48), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

train_datasets = datasets.ImageFolder(os.path.join(data_dir, 'train'), train_transforms) val_datasets = datasets.ImageFolder(os.path.join(data_dir, 'test'), val_transforms) len_train = len(train_datasets) len_val = len(val_datasets)

train_loaders = torch.utils.data.DataLoader( dataset=train_datasets, batch_size=batch_size, shuffle=True, num_workers=0 )

val_loaders = torch.utils.data.DataLoader( dataset=val_datasets, batch_size=batch_size, shuffle=True, num_workers=0 )

criterion = nn.CrossEntropyLoss() optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.8) exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=200, gamma=0.1)

writer = SummaryWriter("runs") best_acc = 0.0

for epoch in range(num_epochs):

start = time.time()

train_loss, train_acc = train(train_loaders, model, criterion, optimizer_ft, device, len_train, batch_size)

exp_lr_scheduler.step()
val_loss, val_acc = val(val_loaders, model, criterion, device, len_val, batch_size)

if val_acc > best_acc:
    best_acc = val_acc
    torch.save(model.state_dict(), './best.pth')
writer.add_scalar('trainloss', train_loss, epoch)
writer.add_scalar('trainacc', train_acc, epoch)
writer.add_scalar('valloss', val_loss, epoch)
writer.add_scalar('valacc', val_acc, epoch)
end = time.time()
print('[{}/{}]: train_loss:{:.3f}, train_acc:{:.3f},eval_loss:{:.3f}, eval_acc:{:.3f},  time:{:.3f}'.format(epoch+1,num_epochs,train_loss,train_acc,val_loss,val_acc,end-start))

writer.close()

PyTorch图像分类模型训练:使用simpleconv网络和Dropout

原文地址: https://www.cveoy.top/t/topic/R2w 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录