pytorch训练一个简单的densenet二分类分类器

 需要设定训练集和验证集：

训练集，验证集里的文件夹名称就是分类的标签：

我放了一些工业品表面的疙瘩和凹坑缺陷图片

train.py如下，batch_size，训练集验证集路径，模型保存路径需要自己手动设置一下：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets, models
from tqdm import tqdm
import os

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据增强和标准化
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 假设输入尺寸为224x224
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

"""---------------------------------------"""
"""---------训练和验证数据集----------------"""
train_dir = 'data/train'  # 修改为你训练数据的路径
val_dir = 'data/val'      # 修改为你验证数据的路径

train_dataset = datasets.ImageFolder(train_dir, transform=transform)
val_dataset = datasets.ImageFolder(val_dir, transform=transform)

"""---------------------------------------"""
"""--------------epoch num----------------"""
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False, num_workers=2)

# 定义模型：DenseNet模型（没有预训练权重）
model = models.densenet121(weights=None)