PyTorch深度学习框架之多分类交叉熵实现图像分类

目录一、自定义小CNN实现手机分类1、代码示例2、代码解析一、自定义小CNN实现手机分类1、代码示例适合苹果/华为/小米 3分类手机识别你可以直接改类别数适配你的任务importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimporttorch.optimasoptimfromtorch.utils.dataimportDataset,DataLoaderfromtorchvisionimporttransformsfromPILimportImageimportos# --------------------------# 第一步定义自定义小CNN 就是我们说的3-4层卷积# 结构输入(3, 224, 224) → 卷积1 → ReLU → 池化 → 卷积2 → ReLU → 池化 → 卷积3 → ReLU → 池化 → 全连接1 → ReLU → 全连接输出# --------------------------classSmallPhoneCNN(nn.Module):def__init__(self,num_classes3): num_classes: 你要分的手机类别数默认3类苹果/华为/小米 super(SmallPhoneCNN,self).__init__()# 1. 卷积层部分3层卷积符合你说的3-4层参数量极小# 输入是RGB图片3通道所以第一层in_channels3self.conv1nn.Conv2d(in_channels3,# 输入RGB三通道out_channels16,# 提取16种特征小任务足够用kernel_size3,# 3×3卷积核最常用padding1# 填充保证卷积后大小不变)self.pool1nn.MaxPool2d(kernel_size2,stride2)# 池化把尺寸缩小一半self.conv2nn.Conv2d(16,32,kernel_size3,padding1)# 输入16通道输出32种特征self.pool2nn.MaxPool2d(2,2)self.conv3nn.Conv2d(32,64,kernel_size3,padding1)# 输入32通道输出64种特征self.pool3nn.MaxPool2d(2,2)# 2. 全连接层部分把卷积提取的特征转成分类输出# 计算一下输入224×224经过3次池化每次缩小一半224 → 112 → 56 → 28# 所以最后特征大小是 64通道 × 28 × 28 64*28*28 50176self.fc1nn.Linear(64*28*28,128)# 卷积特征转128维隐藏向量self.fc2nn.Linear(128,num_classes)# 最后输出128维转成类别数就是我们要的结果# 3. 按照规则做初始化ReLU用Kaiming初始化self._initialize_weights()def_initialize_weights(self):# 所有卷积和全连接都用Kaiming初始化符合我们之前说的规则forminself.modules():ifisinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,modefan_out,nonlinearityrelu)ifm.biasisnotNone:nn.init.zeros_(m.bias)elifisinstance(m,nn.Linear):nn.init.kaiming_normal_(m.weight)nn.init.zeros_(m.bias)# 前向传播定义数据怎么走defforward(self,x):# x形状[batch_size, 3, 224, 224] → 一批图片3通道大小224×224xself.pool1(F.relu(self.conv1(x)))# 第一层卷积激活池化 → 输出[batch,16,112,112]xself.pool2(F.relu(self.conv2(x)))# 第二层 → 输出[batch,32,56,56]xself.pool3(F.relu(self.conv3(x)))# 第三层 → 输出[batch,64,28,28]# 把四维特征拉成一维给全连接层[batch,64,28,28] → [batch, 64*28*28]xx.flatten(1)xF.relu(self.fc1(x))# 全连接第一层激活xself.fc2(x)# 最后一层输出logits不用加softmax训练时候交叉熵会处理returnx# --------------------------# 第二步自定义数据集加载你的手机图片# 数据存放要求和之前一样每个类别一个文件夹# ./phone_data/train/# ├── apple/ (苹果手机图片)# ├── huawei/ (华为手机图片)# └── xiaomi/ (小米手机图片)# ./phone_data/test/ 同理放测试集# --------------------------classPhoneDataset(Dataset):def__init__(self,data_root,transformNone):self.data_rootdata_root self.transformtransform# 读取所有类别生成映射self.class_namessorted(os.listdir(data_root))self.class_to_idx{cls:ifori,clsinenumerate(self.class_names)}# 收集所有图片标签self.img_paths[]forcls_nameinself.class_names:cls_diros.path.join(data_root,cls_name)forimg_nameinos.listdir(cls_dir):self.img_paths.append((os.path.join(cls_dir,img_name),self.class_to_idx[cls_name]))def__len__(self):returnlen(self.img_paths)def__getitem__(self,idx):img_path,labelself.img_paths[idx]imgImage.open(img_path).convert(RGB)ifself.transform:imgself.transform(img)returnimg,label# --------------------------# 第三步训练配置训练流程# --------------------------if__name____main__:# 超参数NUM_CLASSES3# 改成你自己的类别数BATCH_SIZE8EPOCHS15LR1e-3# 数据预处理train_transformtransforms.Compose([transforms.Resize((224,224)),# 直接resize到224小模型不用复杂增强transforms.RandomHorizontalFlip(),# 随机翻转增加数据多样性transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])test_transformtransforms.Compose([transforms.Resize((224,224)),transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# 加载数据集train_datasetPhoneDataset(./phone_data/train,transformtrain_transform)test_datasetPhoneDataset(./phone_data/test,transformtest_transform)train_loaderDataLoader(train_dataset,batch_sizeBATCH_SIZE,shuffleTrue)test_loaderDataLoader(test_dataset,batch_sizeBATCH_SIZE,shuffleFalse)class_namestrain_dataset.class_namesprint(f分类类别{class_names}总训练样本{len(train_dataset)})# 初始化模型、损失、优化器modelSmallPhoneCNN(num_classesNUM_CLASSES)devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)modelmodel.to(device)# 多分类交叉熵和之前讲的完全一样不用改criterionnn.CrossEntropyLoss()optimizeroptim.Adam(model.parameters(),lrLR)# 打印模型参数量看看有多小total_paramssum(p.numel()forpinmodel.parameters())print(f模型总参数量{total_params/1000:.1f}K ≈{total_params/1000000:.2f}M)# 输出大概6.5M比MobileNet还小CPU都能轻松跑# --------------------------# 开始训练# --------------------------forepochinrange(EPOCHS):model.train()train_loss0.0forinputs,labelsintrain_loader:inputs,labelsinputs.to(device),labels.to(device)# 前向算损失outputsmodel(inputs)losscriterion(outputs,labels)# 反向更新optimizer.zero_grad()loss.backward()optimizer.step()train_lossloss.item()*inputs.size(0)# 计算平均损失train_losstrain_loss/len(train_dataset)# 测试集评估准确率model.eval()correct0withtorch.no_grad():forinputs,labelsintest_loader:inputs,labelsinputs.to(device),labels.to(device)outputsmodel(inputs)_,predstorch.max(outputs,1)correcttorch.sum(predslabels.data)acccorrect.double()/len(test_dataset)print(fEpoch [{epoch1}/{EPOCHS}] 训练损失{train_loss:.4f}测试准确率{acc:.4f})# --------------------------# 第四步单张图片预测测试# --------------------------defpredict_single_image(img_path):# 1. model.eval()PyTorch的nn.Module自带方法切换模型到推理模式不用自己写model.eval()# 2. Image.open()PIL库自带方法打开图片不用自己写# 3. .convert(RGB)PIL自带把图片转成三通道RGB处理透明图/灰度图自带的imgImage.open(img_path).convert(RGB)# 4. test_transform我们自己定义的数据预处理就是几个torchvision自带变换的组合所有变换都是自带的我们只需要组合不用自己实现# 5. .unsqueeze(0)PyTorch张量自带方法给张量加一个batch维度自带的# 6. .to(device)PyTorch张量自带方法把张量放到GPU/CPU自带的img_tensortest_transform(img).unsqueeze(0).to(device)# 7. torch.no_grad()PyTorch自带上下文管理器关闭梯度计算自带的withtorch.no_grad():outputsmodel(img_tensor)# model本身我们已经定义好了调用也是自带的# 8. F.softmaxtorch.nn.functional自带函数自带的不需要实现probsF.softmax(outputs,dim1)# 9. torch.maxPyTorch自带函数找最大值和索引自带的top_prob,top_idxtorch.max(probs,dim1)# 10. top_idx[0].item()PyTorch张量自带方法把张量里的单个值转成Python数字自带的pred_clsclass_names[top_idx[0].item()]confidencetop_prob[0].item()print(f\n预测结果输入图片 {img_path})print(f分类结果{pred_cls}置信度{confidence:.4f})returnpred_cls,confidence# 替换成你自己的测试图片predict_single_image(./test_apple.jpg)2、代码解析QSmallPhoneCNN卷积层参数是怎么来的A待完善。。。。。。