90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
import argparse
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torchvision import datasets, transforms
|
|
from torch.utils.data import DataLoader
|
|
from class2.test.mobilenetv3 import MobileNetV3_Large
|
|
|
|
def train(args):
|
|
# 数据增强 & 预处理
|
|
transform_train = transforms.Compose([
|
|
transforms.Resize((224, 224)),
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225])
|
|
])
|
|
transform_val = transforms.Compose([
|
|
transforms.Resize((224, 224)),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
|
std=[0.229, 0.224, 0.225])
|
|
])
|
|
|
|
# 加载数据集
|
|
train_dataset = datasets.ImageFolder(root=f"{args.data}/train", transform=transform_train)
|
|
val_dataset = datasets.ImageFolder(root=f"{args.data}/val", transform=transform_val)
|
|
|
|
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
|
|
val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
|
|
|
|
# 模型
|
|
model = MobileNetV3_Large(num_classes=2)
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = model.to(device)
|
|
|
|
# 损失函数 & 优化器
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=args.lr)
|
|
|
|
# 训练
|
|
for epoch in range(args.epochs):
|
|
model.train()
|
|
running_loss = 0.0
|
|
correct, total = 0, 0
|
|
|
|
for images, labels in train_loader:
|
|
images, labels = images.to(device), labels.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
outputs = model(images)
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
running_loss += loss.item()
|
|
_, predicted = outputs.max(1)
|
|
total += labels.size(0)
|
|
correct += predicted.eq(labels).sum().item()
|
|
|
|
train_acc = 100. * correct / total
|
|
print(f"[Epoch {epoch+1}/{args.epochs}] Loss: {running_loss/len(train_loader):.4f} | Train Acc: {train_acc:.2f}%")
|
|
|
|
# 验证
|
|
model.eval()
|
|
correct, total = 0, 0
|
|
with torch.no_grad():
|
|
for images, labels in val_loader:
|
|
images, labels = images.to(device), labels.to(device)
|
|
outputs = model(images)
|
|
_, predicted = outputs.max(1)
|
|
total += labels.size(0)
|
|
correct += predicted.eq(labels).sum().item()
|
|
|
|
val_acc = 100. * correct / total
|
|
print(f"Validation Acc: {val_acc:.2f}%")
|
|
|
|
# 保存模型
|
|
torch.save(model.state_dict(), "mobilenetv3_binary.pth")
|
|
print("训练完成，模型已保存到 mobilenetv3_binary.pth")
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--data", type=str, default="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classdata", help="数据集路径")
|
|
parser.add_argument("--epochs", type=int, default=100, help="训练轮数")
|
|
parser.add_argument("--batch_size", type=int, default=32, help="批大小")
|
|
parser.add_argument("--lr", type=float, default=1e-3, help="学习率")
|
|
args = parser.parse_args()
|
|
train(args)
|