复制代码
import torch
import torch.nn as nn
from torchvision import transforms, datasets
import os,PIL,pathlib,warnings
warnings.filterwarnings("ignore") #忽略警告信息
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
import matplotlib.pyplot as plt
from PIL import Image
import os
image_folder = './data/data/2Mild/' # 指定图像文件夹路径
# 获取文件夹中的所有图像文件
image_files = [f for f in os.listdir(image_folder) if f.endswith((".jpg", ".png", ".jpeg"))]
fig, axes = plt.subplots(3, 8, figsize=(16, 6)) # 创建Matplotlib图像
# 使用列表推导式加载和显示图像
for ax, img_file in zip(axes.flat, image_files):
img_path = os.path.join(image_folder, img_file)
img = Image.open(img_path)
ax.imshow(img)
ax.axis('off')
# 显示图像
plt.tight_layout()
plt.show()
data_dir = './data/data/'
# 关于transforms.Compose的更多介绍可以参考:https://blog.csdn.net/qq_38251616/article/details/124878863
train_transforms = transforms.Compose([
transforms.Resize([224, 224]), # 将输入图片resize成统一尺寸
transforms.ToTensor(), # 将PIL Image或numpy.ndarray转换为tensor,并归一化到[0,1]之间
transforms.Normalize( # 标准化处理->转换为标准正太分布(高斯分布),使模型更容易收敛
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # 其中 mean与std是从数据集中随机抽样计算得到的
])
total_data = datasets.ImageFolder(data_dir, transform=train_transforms)
total_data
total_data.class_to_idx
# 按8:2比例划分训练集和测试集
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
# 随机拆分数据集
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])
batch_size = 4
# 创建训练集数据加载器(打乱数据)
train_dl = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True
)
# 创建测试集数据加载器(不打乱数据)
test_dl = torch.utils.data.DataLoader(
test_dataset,
batch_size=batch_size
)
# 查看数据加载器的输出格式
for X, y in test_dl:
print("Shape of X [N, C, H, W]: ", X.shape)
print("Shape of y: ", y.shape, y.dtype)
break
复制代码
import torch
import torch.nn as nn
# Same Padding:自动计算卷积的填充大小
def autopad(k, p=None): # kernel, padding
# pad to 'same'
if p is None:
if isinstance(k, int):
p = k // 2 # 如果k是整数,填充为k的一半
else:
p = [x // 2 for x in k] # 如果k是列表,每个元素取一半
return p
# Identity Block(残差块,输入输出维度一致)
class IdentityBlock(nn.Module):
def __init__(self, in_channel, kernel_size, filters):
super(IdentityBlock, self).__init__()
filters1, filters2, filters3 = filters
self.conv1 = nn.Sequential(
nn.Conv2d(in_channel, filters1, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(filters1),
nn.ReLU(True)
)
self.conv2 = nn.Sequential(
nn.Conv2d(filters1, filters2, kernel_size, stride=1, padding=autopad(kernel_size), bias=False),
nn.BatchNorm2d(filters2),
nn.ReLU(True)
)
self.conv3 = nn.Sequential(
nn.Conv2d(filters2, filters3, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(filters3)
)
self.relu = nn.ReLU(True)
def forward(self, x):
x1 = self.conv1(x)
x1 = self.conv2(x1)
x1 = self.conv3(x1)
x = x1 + x # 残差连接
x = self.relu(x)
return x
# Conv Block(残差块,输入输出维度不一致,需卷积调整维度)
class ConvBlock(nn.Module):
def __init__(self, in_channel, kernel_size, filters, stride=2):
super(ConvBlock, self).__init__()
filters1, filters2, filters3 = filters
self.conv1 = nn.Sequential(
nn.Conv2d(in_channel, filters1, 1, stride=stride, padding=0, bias=False),
nn.BatchNorm2d(filters1),
nn.ReLU(True)
)
self.conv2 = nn.Sequential(
nn.Conv2d(filters1, filters2, kernel_size, stride=1, padding=autopad(kernel_size), bias=False),
nn.BatchNorm2d(filters2),
nn.ReLU(True)
)
self.conv3 = nn.Sequential(
nn.Conv2d(filters2, filters3, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(filters3)
)
self.conv4 = nn.Sequential(
nn.Conv2d(in_channel, filters3, 1, stride=stride, padding=0, bias=False),
nn.BatchNorm2d(filters3)
)
self.relu = nn.ReLU(True)
def forward(self, x):
x1 = self.conv1(x)
x1 = self.conv2(x1)
x1 = self.conv3(x1)
x2 = self.conv4(x) # 调整输入维度以匹配输出
x = x1 + x2 # 残差连接
x = self.relu(x)
return x
# 构建ResNet-50网络
class ResNet50(nn.Module):
def __init__(self, classes=1000):
super(ResNet50, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False, padding_mode='zeros'),
nn.BatchNorm2d(64),
nn.ReLU(True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)
self.conv2 = nn.Sequential(
ConvBlock(64, 3, [64, 64, 256], stride=1),
IdentityBlock(256, 3, [64, 64, 256]),
IdentityBlock(256, 3, [64, 64, 256])
)
self.conv3 = nn.Sequential(
ConvBlock(256, 3, [128, 128, 512]),
IdentityBlock(512, 3, [128, 128, 512]),
IdentityBlock(512, 3, [128, 128, 512]),
IdentityBlock(512, 3, [128, 128, 512])
)
self.conv4 = nn.Sequential(
ConvBlock(512, 3, [256, 256, 1024]),
IdentityBlock(1024, 3, [256, 256, 1024]),
IdentityBlock(1024, 3, [256, 256, 1024]),
IdentityBlock(1024, 3, [256, 256, 1024]),
IdentityBlock(1024, 3, [256, 256, 1024]),
IdentityBlock(1024, 3, [256, 256, 1024])
)
self.conv5 = nn.Sequential(
ConvBlock(1024, 3, [512, 512, 2048]),
IdentityBlock(2048, 3, [512, 512, 2048]),
IdentityBlock(2048, 3, [512, 512, 2048])
)
self.pool = nn.AvgPool2d(kernel_size=7, stride=7, padding=0)
self.fc = nn.Linear(2048, classes) # 这里classes设为3,对应分类任务
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.conv5(x)
x = self.pool(x)
x = torch.flatten(x, start_dim=1)
x = self.fc(x)
return x
# 实例化模型并移动到设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResNet50(classes=3).to(device)
# 统计模型参数数量以及其他指标
import torchsummary as summary
summary.summary(model, (3, 224, 224))
复制代码
# 训练循环
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) # 训练集的大小
num_batches = len(dataloader) # 批次数目(size/batch_size,向上取整)
train_loss, train_acc = 0, 0 # 初始化训练损失和正确率
for X, y in dataloader: # 获取图片及其标签
X, y = X.to(device), y.to(device)
# 计算预测误差
pred = model(X) # 网络输出
loss = loss_fn(pred, y) # 计算网络输出和真实值之间的差距
# 反向传播
optimizer.zero_grad() # grad属性归零
loss.backward() # 反向传播
optimizer.step() # 每一步自动更新
# 记录acc和loss
train_acc += (pred.argmax(1) == y).type(torch.float).sum().item()
train_loss += loss.item()
train_acc /= size # 计算训练集整体正确率
train_loss /= num_batches # 计算训练集平均损失
return train_acc, train_loss
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset) # 测试集的大小
num_batches = len(dataloader) # 批次数目(size/batch_size,向上取整)
test_loss, test_acc = 0, 0
# 当不进行训练时,停止梯度更新,节省计算内存消耗
with torch.no_grad():
for imgs, target in dataloader:
imgs, target = imgs.to(device), target.to(device)
# 计算loss
target_pred = model(imgs)
loss = loss_fn(target_pred, target)
test_loss += loss.item()
test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()
test_acc /= size
test_loss /= num_batches
return test_acc, test_loss
import copy
import torch
import torch.nn as nn
import torch.optim as optim
# 初始化优化器与损失函数
optimizer = optim.AdamW(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss() # 创建损失函数
epochs = 10 # 训练轮数
# 初始化指标记录列表
train_loss = []
train_acc = []
test_loss = []
test_acc = []
best_acc = 0 # 设置最佳准确率,作为保存最佳模型的指标
for epoch in range(epochs):
# 训练阶段
model.train() # 开启训练模式(启用Dropout、BatchNorm等层的训练行为)
epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, optimizer)
# 测试阶段
model.eval() # 开启评估模式(禁用Dropout、固定BatchNorm等层的参数)
epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)
# 保存最佳模型
if epoch_test_acc > best_acc:
best_acc = epoch_test_acc
best_model = copy.deepcopy(model) # 深拷贝当前最佳模型
# 记录训练/测试指标
train_acc.append(epoch_train_acc)
train_loss.append(epoch_train_loss)
test_acc.append(epoch_test_acc)
test_loss.append(epoch_test_loss)
# 获取当前学习率
lr = optimizer.state_dict()['param_groups'][0]['lr']
# 打印当前轮次的指标
template = ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}')
print(template.format(epoch+1,
epoch_train_acc*100,
epoch_train_loss,
epoch_test_acc*100,
epoch_test_loss,
lr))
# 保存最佳模型到文件
PATH = './best_model.pth' # 保存的参数文件名
torch.save(best_model.state_dict(), PATH) # 保存模型的参数状态字典
print('Done')
import matplotlib.pyplot as plt
# 隐藏警告
import warnings
warnings.filterwarnings("ignore") # 忽略警告信息
# 配置Matplotlib显示(解决中文/负号显示问题)
plt.rcParams['font.sans-serif'] = ['SimHei'] # 正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False # 正常显示负号
plt.rcParams['figure.dpi'] = 100 # 设置图像分辨率
from datetime import datetime
current_time = datetime.now() # 获取当前时间
epochs_range = range(epochs)
# 创建画布并绘制子图
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
# 绘制训练/测试准确率曲线
plt.plot(epochs_range, train_acc, label='Training Accuracy')
plt.plot(epochs_range, test_acc, label='Test Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
plt.xlabel(current_time) # 横轴标注当前时间(打卡用)
plt.subplot(1, 2, 2)
# 绘制训练/测试损失曲线
plt.plot(epochs_range, train_loss, label='Training Loss')
plt.plot(epochs_range, test_loss, label='Test Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()