mobilenet图像分类模型代码含训练和预测

### 使用 MobileNet 进行图像分类的训练和预测 #### 创建 MobileNet 模型 为了构建一个用于图像分类的任务，可以使用预定义好的 `MobileNet` 架构。对于 PyTorch 用户来说，可以通过 torchvision 库方便地获取该架构。 ```python import torch from torchvision import models, transforms from torch.nn import functional as F ``` #### 修改模型结构适应特定任务需求 当处理自定义的数据集时，通常需要调整最后一层以匹配目标类别的数目。这里展示了如何更改 MobileNetV2 的分类器部分来适配新的类别数[^2]： ```python model = models.mobilenet_v2(pretrained=True) # 冻结除了 classifier 外的所有参数 for param in model.parameters(): param.requires_grad = False n_inputs = model.classifier[1].in_features last_layer = torch.nn.Linear(n_inputs, num_classes) # 假设num_classes是你数据集中类的数量 model.classifier[1] = last_layer ``` #### 设置训练选项与优化策略 设置合适的超参数对于获得良好的泛化能力至关重要。这包括但不限于批量大小(batch size)，初始学习率(learning rate)等配置项。此外，在此阶段还可以指定不同的学习速率给不同层次的权重更新过程。 ```python criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam([ {'params': model.features.parameters(), 'lr': 0.001}, {'params': model.classifier.parameters(), 'lr': 0.01} ]) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) ``` #### 数据加载与预处理 准备输入到神经网络之前的数据非常重要；一般会涉及到标准化(normalization), 裁剪(cropping), 翻转(flipping)等一系列操作。这些变换有助于增强模型的表现力以及防止过拟合现象的发生。 ```python data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), 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]) ]), } image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in ['train', 'val']} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes ``` #### 开始训练循环 一旦所有的准备工作都完成了，就可以进入实际的训练环节了。在这个过程中，不断地迭代整个数据集，并通过反向传播算法最小化损失函数值。 ```python device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") def train_model(model, criterion, optimizer, scheduler, num_epochs=25): since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval() running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # zero the parameter gradients optimizer.zero_grad() # forward with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # backward + optimize only if in training phase if phase == 'train': loss.backward() optimizer.step() # statistics running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) if phase == 'train': scheduler.step() epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print('{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) # deep copy the model if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) # load best model weights model.load_state_dict(best_model_wts) return model ``` #### 执行预测 完成训练后，利用已保存的最佳模型来进行新样本上的推断是非常简单的一步。只需要将待测试图片按照同样的方式转换成张量形式即可得到其所属类别概率分布情况。 ```python def predict_image(img_path, model): img = Image.open(img_path).convert('RGB') preprocess = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input_tensor = preprocess(img) input_batch = input_tensor.unsqueeze(0) out = model(input_batch) probabilities = F.softmax(out, dim=1)[0] top_prob, top_class = torch.topk(probabilities, k=1) return class_names[top_class.cpu().numpy()[0]], float(top_prob) ```