PyTorch图像分类模型训练:使用simpleconv网络和Dropout
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()
原文地址: https://www.cveoy.top/t/topic/R2w 著作权归作者所有。请勿转载和采集!