import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, accuracy_score, recall_score, f1_score, precision_score
import seaborn as sns
import os
# 定义基础的残差块(BasicBlock)
class BasicBlock(nn.Module):
expansion = 1 # 残差块的扩展系数,用于调整通道数
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
# 第一个卷积层,3x3卷积核,步长为stride,填充为1
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes) # 批归一化层
# 第二个卷积层,3x3卷积核,步长为1,填充为1
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes) # 批归一化层
# 残差连接中的shortcut,根据stride和输入输出通道数是否一致来决定是否需要进行通道变换
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion * planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion * planes)
)
def forward(self, x):
identity = x # 保存输入的identity
out = F.relu(self.bn1(self.conv1(x))) # 第一层卷积、批归一化和ReLU激活
out = self.bn2(self.conv2(out)) # 第二层卷积和批归一化
out += self.shortcut(identity) # 残差连接
out = F.relu(out) # 最后再通过ReLU激活
return out
# 定义ResNet34模型
class ResNet34(nn.Module):
def __init__(self, block, num_blocks, num_classes=10):
super(ResNet34, self).__init__()
self.in_planes = 64 # 输入通道数
# 第一层:7x7卷积核,步长为2,填充为3
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64) # 批归一化层
# 四个阶段的残差块组成的层
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
# 全连接层,输出为num_classes个类别
self.linear = nn.Linear(512 * block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1) # 第一个残差块的步长为stride,后面的都为1
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x))) # 第一层卷积、批归一化和ReLU激活
out = F.max_pool2d(out, kernel_size=3, stride=2, padding=1) # 7x7卷积后添加一个3x3最大池化层
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4) # 全局平均池化
out = out.view(out.size(0), -1) # 展平
out = self.linear(out) # 全连接层
return out
# 参数设置
num_classes = 10
batch_size = 16
num_epochs = 20
lr = 0.0001
# 数据增强
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
# 载入数据集
dataset = ImageFolder('D:/实习/dataset1', transform=transform)
torch.manual_seed(42)
# 划分训练集、验证集、测试集
train_size = int(0.7 * len(dataset))
val_size = int(0.2 * len(dataset))
test_size = len(dataset) - train_size - val_size
train_set, val_set, test_set = random_split(dataset, [train_size, val_size, test_size], generator=torch.Generator().manual_seed(42))
# 创建数据加载器
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_set, batch_size=batch_size)
test_loader = DataLoader(test_set, batch_size=batch_size)
# 定义模型所需的参数
block = BasicBlock
num_blocks = [3, 4, 6, 3]
# 初始化模型
model = ResNet34(block, num_blocks, num_classes=num_classes)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
# 设置设备(CPU或GPU)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 将模型移至设备
model.to(device)
# 学习率调度器
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
# 训练过程中收集的准确率和损失数据
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []
# 保存最佳模型的参数
best_val_accuracy = 0.0
best_model_wts = None
# 训练
for epoch in range(num_epochs):
running_loss = 0.0
correct = 0
total = 0
model.train()
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
# 正向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted.to(device) == labels.to(device)).sum().item()
# 打印训练集的损失率和准确率
epoch_loss = running_loss / len(train_set)
epoch_accuracy = 100 * correct / total
# 收集训练和验证集的准确率和损失数据
train_losses.append(epoch_loss)
train_accuracies.append(epoch_accuracy)
print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}, Train Accuracy: {epoch_accuracy:.2f}%")
# 验证
model.eval()
correct = 0
total = 0
total_loss = 0.0
with torch.no_grad():
for inputs, labels in val_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted.to(device) == labels.to(device)).sum().item()
loss = criterion(outputs, labels)
total_loss += loss.item()
val_accuracy = 100 * correct / total
val_loss = total_loss / len(val_loader)
val_losses.append(val_loss)
val_accuracies.append(val_accuracy)
print(f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
# 保存最佳模型
if val_accuracy > best_val_accuracy:
best_val_accuracy = val_accuracy
best_model_wts = model.state_dict()
# 更新学习率
scheduler.step()
# 加载最佳模型
model.load_state_dict(best_model_wts)
# 打印最佳模型的参数
print("Best model parameters:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 确保保存模型的目录存在
save_dir = 'D:/实习/models'
os.makedirs(save_dir, exist_ok=True)
# 保存最佳模型
model_path = os.path.join(save_dir, 'best_model.pth')
torch.save(model.state_dict(), model_path)
# 测试
y_pred = []
y_true = []
model.eval()
test_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
y_pred.extend(predicted.tolist())
y_true.extend(labels.tolist())
total += labels.size(0)
correct += (predicted == labels).sum().item()
test_loss += criterion(outputs, labels)
test_accuracy = (100 * correct / total)
test_loss = test_loss / len(test_loader)
print(f'Test Loss: {test_loss}, Test Accuracy: {test_accuracy:.2f}%')
# 计算混淆矩阵
conf_matrix = confusion_matrix(y_true, y_pred)
# 可视化混淆矩阵
plt.figure(figsize=(10, 8))
sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', xticklabels=range(num_classes), yticklabels=range(num_classes))
plt.title('Confusion Matrix')
plt.xlabel('Predicted')
plt.ylabel('True')
plt.savefig('D:/实习/result/Confusion_Matrix.png')
# 打印其他指标
print(f"Accuracy: {accuracy_score(y_true, y_pred)}")
print(f"Recall: {recall_score(y_true, y_pred, average='macro')}")
print(f"F1 Score: {f1_score(y_true, y_pred, average='macro')}")
print(f"Precision: {precision_score(y_true, y_pred, average='macro')}")
# 绘制训练和验证集的准确率图表
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(val_accuracies, label='Validation Accuracy')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.grid(True)
# 绘制训练和验证集的损失图表
plt.subplot(1, 2, 2)
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Validation Loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)
plt.savefig('D:/实习/t-v-l/training_validation_loss_accuracy.png')
怎么修改