pythonproject/qrtrain.py

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='')