222 lines
8.4 KiB
Python
222 lines
8.4 KiB
Python
import os
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.optim.lr_scheduler import StepLR
|
||
from torchvision import models, transforms
|
||
from torch.utils.data import DataLoader
|
||
from torchvision.datasets import ImageFolder
|
||
from torch.optim import Adam
|
||
from tqdm import tqdm
|
||
from torchvision.transforms import functional
|
||
from FocalLoss import FocalLoss
|
||
from torchvision.models import ResNet18_Weights
|
||
from torch.cuda.amp import GradScaler, autocast
|
||
|
||
# 加载预训练的ResNet-18模型
|
||
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
|
||
|
||
# 将模型移到GPU上（如果有）
|
||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||
model = model.to(device)
|
||
|
||
# 替换最后一层以适应二分类任务
|
||
num_features = model.fc.in_features
|
||
model.fc = nn.Linear(num_features, 2) # 假设是二分类问题
|
||
model.fc = model.fc.to(device)
|
||
|
||
# 冻结所有层
|
||
for param in model.parameters():
|
||
param.requires_grad = False
|
||
|
||
# 解冻最后一层，用于微调
|
||
# for param in model.fc.parameters():
|
||
# param.requires_grad = True
|
||
|
||
# 解冻最后几层
|
||
for name, param in model.named_parameters():
|
||
if "layer4" in name or "fc" in name: # 只解冻最后几层
|
||
param.requires_grad = True
|
||
|
||
def random_brightness(image):
|
||
"""
|
||
随机调整图像的亮度。
|
||
|
||
该函数通过生成一个随机的亮度因子，然后应用到图像上，以增加或减少图像的亮度。这种方法常用于数据增强，
|
||
旨在帮助模型更好地泛化，因为它教模型在不同亮度条件下处理图像。
|
||
|
||
参数:
|
||
image (Tensor): 输入的图像Tensor。
|
||
|
||
返回:
|
||
Tensor: 亮度经过随机调整后的图像Tensor。
|
||
"""
|
||
# 生成一个随机的亮度因子
|
||
brightness_factor = torch.rand(1).item()
|
||
# 调整图像的亮度
|
||
return functional.adjust_brightness(image, brightness_factor)
|
||
|
||
|
||
def random_contrast(image):
|
||
"""
|
||
随机调整图像的对比度。
|
||
|
||
该函数通过生成一个随机的对比度因子，然后应用到图像上，以增加或减少图像的对比度。
|
||
对比度因子是在0.5到1.5之间随机生成的，其中0.5是最低对比度，1.5是最高对比度。
|
||
|
||
参数:
|
||
image (Tensor): 输入的图像Tensor。
|
||
|
||
返回:
|
||
Tensor: 对比度经过随机调整后的图像Tensor。
|
||
"""
|
||
# 生成一个随机的对比度因子，其值在0.5到1.5之间。
|
||
contrast_factor = torch.rand(1).item() + 0.5
|
||
# 使用生成的对比度因子调整图像的对比度。
|
||
return functional.adjust_contrast(image, contrast_factor)
|
||
|
||
|
||
def random_hue(image):
|
||
"""
|
||
随机调整图像的色相。
|
||
|
||
色相的变化是由一个随机生成的因子决定的，这个因子是在-0.05到0.05之间均匀分布的。
|
||
这个函数的目的是为图像数据增强，使得模型在训练过程中看到更多变体，从而提高泛化能力。
|
||
|
||
参数:
|
||
image (Tensor): 输入的图像Tensor。
|
||
|
||
返回:
|
||
Tensor: 色相被随机调整后的图像Tensor。
|
||
"""
|
||
# 生成一个随机的色相调整因子，其值在-0.05到0.05之间。
|
||
hue_factor = (torch.rand(1).item() - 0.5) / 10
|
||
# 使用调整的色相因子来改变图像的色相。
|
||
return functional.adjust_hue(image, hue_factor)
|
||
|
||
|
||
def add_gaussian_noise(image):
|
||
"""
|
||
向图像数据添加高斯噪声。
|
||
|
||
该函数的目的是为了模拟现实世界中图像经常受到的各种噪声干扰。通过在图像中添加高斯噪声，
|
||
可以使模型在训练过程中更加鲁棒，提高其在真实场景下的泛化能力。
|
||
|
||
参数:
|
||
image: 输入的图像数据，作为一个torch张量。
|
||
|
||
返回值:
|
||
返回添加了高斯噪声的图像数据，与输入图像数据类型和形状相同。
|
||
|
||
添加的噪声是通过torch.randn_like函数生成的，确保了噪声的分布与图像数据相同，
|
||
并且乘以了一个小的常数0.1来控制噪声的强度。
|
||
"""
|
||
# 生成与图像数据相同形状的高斯噪声张量，并将其与图像数据相加
|
||
return image + torch.randn_like(image) * 0.1
|
||
|
||
|
||
data_transforms = {
|
||
'train': transforms.Compose([
|
||
transforms.RandomResizedCrop(224),
|
||
transforms.RandomHorizontalFlip(),
|
||
transforms.RandomRotation(degrees=15), # 随机旋转图像最多15度
|
||
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1), # 随机改变图像的颜色
|
||
transforms.ToTensor(),
|
||
# transforms.Lambda(lambda x: functional.adjust_brightness(x, brightness_factor=torch.rand(1).item())), # 随机调整亮度
|
||
# transforms.Lambda(lambda x: functional.adjust_contrast(x, contrast_factor=torch.rand(1).item() + 0.5)), # 随机调整对比度
|
||
# transforms.Lambda(lambda x: functional.adjust_hue(x, hue_factor=(torch.rand(1).item()-0.5)/10)), # 随机调整色相
|
||
# transforms.Lambda(lambda x: x + torch.randn_like(x) * 0.1), # 添加随机高斯噪声
|
||
transforms.Lambda(random_brightness),
|
||
transforms.Lambda(random_contrast),
|
||
transforms.Lambda(random_hue),
|
||
transforms.Lambda(add_gaussian_noise),
|
||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||
]),
|
||
'val': transforms.Compose([
|
||
transforms.Resize(256),
|
||
transforms.CenterCrop(224),
|
||
transforms.ToTensor(),
|
||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||
]),
|
||
}
|
||
|
||
# 加载数据集
|
||
data_dir = '/Users/jasonwong/Downloads/二维码测试1/data'
|
||
image_datasets = {x: ImageFolder(os.path.join(data_dir, x), data_transforms[x])
|
||
for x in ['train', 'val']}
|
||
dataloaders = {x: DataLoader(image_datasets[x], batch_size=32, shuffle=True, num_workers=4)
|
||
for x in ['train', 'val']}
|
||
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
|
||
|
||
#设置损失函数和优化器】
|
||
num_normal_samples = 7891
|
||
num_anomaly_samples = 846
|
||
class_sample_counts = [num_normal_samples, num_anomaly_samples] # 替换为实际样本数量
|
||
weights = 1.0 / np.array(class_sample_counts)
|
||
weights /= weights.sum() # 确保权重和为1
|
||
class_weights = torch.FloatTensor(weights).to(device) # 将权重转换为PyTorch张量，并移动到设备上
|
||
# criterion = nn.CrossEntropyLoss()
|
||
# optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
|
||
optimizer = Adam(model.parameters(),
|
||
lr=1e-4, # 学习率
|
||
betas=(0.9, 0.999), # beta_1 和 beta_2
|
||
eps=1e-08) # epsilon
|
||
scheduler = StepLR(optimizer, step_size=7, gamma=0.1)
|
||
|
||
# 初始化GradScaler
|
||
# scaler = GradScaler()
|
||
|
||
#训练模型
|
||
def train_model(model, dataloaders, dataset_sizes, device, num_epochs=10, save_path='path_to_save/full_model'):
|
||
# criterion = FocalLoss(gamma=2, alpha=[0.25, 0.75], reduction='mean')
|
||
# 加载加权交叉熵损失函数
|
||
criterion = nn.CrossEntropyLoss(weight=class_weights)
|
||
|
||
for epoch in range(num_epochs):
|
||
print(f'Starting epoch {epoch + 1}/{num_epochs}')
|
||
|
||
for phase in ['train', 'val']:
|
||
if phase == 'train':
|
||
model.train()
|
||
else:
|
||
model.eval()
|
||
|
||
running_loss = 0.0
|
||
running_corrects = 0
|
||
|
||
for inputs, labels in tqdm(dataloaders[phase], desc=f'Epoch {epoch + 1} {phase.capitalize()}'):
|
||
inputs = inputs.to(device)
|
||
labels = labels.to(device)
|
||
|
||
optimizer.zero_grad()
|
||
|
||
with torch.set_grad_enabled(phase == 'train'):
|
||
outputs = model(inputs)
|
||
_, preds = torch.max(outputs, 1)
|
||
loss = criterion(outputs, labels)
|
||
|
||
if phase == 'train':
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
running_loss += loss.item() * inputs.size(0)
|
||
running_corrects += torch.sum(preds == labels.data)
|
||
|
||
epoch_loss = running_loss / dataset_sizes[phase]
|
||
epoch_acc = running_corrects.double() / dataset_sizes[phase]
|
||
|
||
print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
|
||
|
||
scheduler.step()
|
||
|
||
torch.save({
|
||
'epoch': epoch + 1,
|
||
'model_state_dict': model.state_dict(),
|
||
'optimizer_state_dict': optimizer.state_dict(),
|
||
'loss': epoch_loss,
|
||
}, f'{save_path}_epoch_{epoch + 1}.pth')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
train_model(model, dataloaders, dataset_sizes, device, num_epochs=10, save_path='') |