resnet分类训练

  1. resnet分类器训练
复制代码
import torch
import torchvision
from torchvision import transforms
from torch.utils.data import random_split
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet50

# Define the transformation
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Load the dataset
data = torchvision.datasets.ImageFolder(root=r"D:\train_model\train_data_set", transform=transform)

classes_set = data.classes
# 保存类别信息到 classes.txt
with open('classes.txt', 'w') as f:
    for class_name in classes_set:
        f.write(class_name + '\n')
# Split the data into train and test sets
train_size = int(0.8 * len(data))
test_size = len(data) - train_size
train_data, test_data = random_split(data, [train_size, test_size])

# Optionally, you can load the train and test data into data loaders
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=False)

# Define the model
model = resnet50(pretrained=True)

# Replace the last layer
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, len(classes_set))
# Define the loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Move the model to the device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Define the number of epochs
num_epochs = 10

# Train the model
for epoch in range(num_epochs):
    # Train the model on the training set
    model.train()
    train_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        # Move the data to the device
        inputs = inputs.to(device)
        # inputs = inputs.float()
        labels = labels.to(device)
        # labels = labels.long()

        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward + backward + optimize
        outputs = model(inputs)

        loss = criterion(outputs, labels)

        loss.backward()

        optimizer.step()

        # Update the training loss
        train_loss += loss.item() * inputs.size(0)

    # Evaluate the model on the test set
    model.eval()
    test_loss = 0.0
    test_acc = 0.0
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(test_loader):
            # Move the data to the device
            inputs = inputs.to(device)
            labels = labels.to(device)

            # Forward
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # Update the test loss and accuracy
            test_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            test_acc += torch.sum(preds == labels.data)

    # Print the training and test loss and accuracy
    train_loss /= len(train_data)
    test_loss /= len(test_data)
    test_acc = test_acc.double() / len(test_data)
    print(f"Epoch [{epoch + 1}/{num_epochs}] Train Loss: {train_loss:.4f} Test Loss: {test_loss:.4f} Test Acc: {test_acc:.4f}")

# 保存模型参数
torch.save(model.state_dict(), './model/trained_model.pth')
相关推荐
AI人工智能+1 小时前
炫光活体检测技术:通过光学技术实现高效、安全的身份验证,有效防御多种伪造手段。
人工智能·深度学习·人脸识别·活体检测
SHUIPING_YANG2 小时前
如何让dify分类器更加精准的分类?
人工智能·分类·数据挖掘
东方佑2 小时前
打破常规:“无注意力”神经网络为何依然有效?
人工智能·深度学习·神经网络
Francek Chen3 小时前
【深度学习计算机视觉】03:目标检测和边界框
人工智能·pytorch·深度学习·目标检测·计算机视觉·边界框
九章云极AladdinEdu3 小时前
AI集群全链路监控:从GPU微架构指标到业务Metric关联
人工智能·pytorch·深度学习·架构·开源·gpu算力
惯导马工3 小时前
【论文导读】IDOL: Inertial Deep Orientation-Estimation and Localization
深度学习·算法
爱学习的茄子3 小时前
Function Call:让AI从文本生成走向智能交互的技术革命
前端·深度学习·openai
CoovallyAIHub3 小时前
基于YOLO集成模型的无人机多光谱风电部件缺陷检测
深度学习·算法·计算机视觉
CoovallyAIHub3 小时前
几十个像素的小目标,为何难倒无人机?LCW-YOLO让无人机小目标检测不再卡顿
深度学习·算法·计算机视觉
IMER SIMPLE4 小时前
人工智能-python-深度学习-经典网络模型-LeNets5
人工智能·python·深度学习