pytorch训练mobilenetv1完整程序

PyTorch 中训练 MobileNetV1 的完整程序通常包括以下几个步骤：

1. **导入必要的库**：

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

2. **下载预训练模型并加载**：

model = models.mobilenet_v2(pretrained=True)
num_features = model.fc.in_features  # 获取最后一层（全连接层）的输入特征数

3. **冻结前几层**（如果不需要微调所有层）：

for param in model.parameters():
    param.requires_grad = False
model.fc = torch.nn.Linear(num_features, num_classes)  # 定义新的分类层

4. **设置损失函数和优化器**：

criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

5. **数据处理**：

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

train_dataset = datasets.ImageFolder(train_dir, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

val_dataset = datasets.ImageFolder(val_dir, transform=transform)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

6. **训练循环**：

num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 计算验证集精度
    with torch.no_grad():
        correct = 0
        total = 0
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        val_accuracy = 100 * correct / total
        print(f'Epoch {epoch+1}/{num_epochs}, Validation Accuracy: {val_accuracy}%')

7. **保存模型**：

torch.save(model.state_dict(), 'mobile_net.pth')

在这个例子中，`train_dir` 和 `val_dir` 分别指定了训练集和验证集的目录，`batch_size` 是每次迭代使用的样本数量，`num_classes` 是分类任务的类别数。