path= "/home/ma-user/work/Garbage_Data" import mindspore.dataset as ds import mindspore.dataset.vision as vision def create_dataset(path, batch_size=50, train=True, image_size=224): #dataset = ds.ImageFolderDataset(path, num_parallel_workers=8, class_indexing={"croissants": 0, "dog": 1}) dataset = ds.ImageFolderDataset(path, num_parallel_workers=8, class_indexing={"cardboard": 0, "glass": 1, "metal": 2, "paper": 3, "plastic": 4, "trash": 5}) # 图像增强操作 mean = [0.485 * 255, 0.456 * 255, 0.406 * 255] std = [0.229 * 255, 0.224 * 255, 0.225 * 255] if train: trans = [ vision.RandomCropDecodeResize(image_size, scale=(0.08, 1.0), ratio=(0.75, 1.333)), vision.RandomHorizontalFlip(prob=0.5), vision.Normalize(mean=mean, std=std), vision.HWC2CHW() ] else: trans = [ vision.Decode(), vision.Resize(256), vision.CenterCrop(image_size), vision.Normalize(mean=mean, std=std), vision.HWC2CHW() ] dataset = dataset.map(operations=trans, input_columns="image", num_parallel_workers=8) # 设置batch_size的大小,若最后一次抓取的样本数小于batch_size,则丢弃 dataset = dataset.batch(batch_size, drop_remainder=True) return dataset # 加载训练数据集和验证数据集用于后续的模型训练和验证。 # 加载训练数据集 train_path = "Garbage_Data/train" dataset_train = create_dataset(train_path, train=True) # 加载验证数据集 val_path = "Garbage_Data/val" dataset_val = create_dataset(val_path, train=False) from mindvision.dataset import DownLoad models_url = "https://download.mindspore.cn/vision/classification/mobilenet_v2_1.0_224.ckpt" dl = DownLoad() # 下载预训练模型文件 dl.download_url(models_url) # ### MobileNet V2模型微调 # # 本章使用MobileNet V2的预训练模型进行微调,通过删除MobileNet V2预训练模型中最后一个用于分类的1x1的卷积层的参数,使用狗与牛角包数据集对模型进行重新训练以更新模型参数。 # In[5]: import mindspore.nn as nn import mindspore as ms from mindvision.classification.models import mobilenet_v2 from mindvision.engine.loss import CrossEntropySmooth # 创建模型,其中目标分类数为6,图像输入大小为(224,224) network = mobilenet_v2(num_classes=6, resize=224) # 模型参数存入到param_dict param_dict = ms.load_checkpoint("./mobilenet_v2_1.0_224.ckpt") # 获取mobilenet_v2网络最后一个卷积层的参数名 filter_list = [x.name for x in network.head.classifier.get_parameters()] # 删除预训练模型的最后一个卷积层 def filter_ckpt_parameter(origin_dict, param_filter): for key in list(origin_dict.keys()): for name in param_filter: if name in key: print("Delete parameter from checkpoint: ", key) del origin_dict[key] break filter_ckpt_parameter(param_dict, filter_list) # 加载预训练模型参数作为网络初始化权重 ms.load_param_into_net(network, param_dict) # 定义优化器 network_opt = nn.Momentum(params=network.trainable_params(), learning_rate=0.01, momentum=0.9) # 定义损失函数 network_loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=0.1, classes_num=2) # 定义评价指标 metrics = {"Accuracy": nn.Accuracy()} # 初始化模型 model = ms.Model(network, loss_fn=network_loss, optimizer=network_opt, metrics=metrics) # # ### 模型训练与评估 # # 训练并评估网络,使用MindSpore Vision中的`mindvision.engine.callback.ValAccMonitor`接口打印训练的损失值和评估精度,且在训练完成后,保存评估精度最高的CKPT文件`best.ckpt`在当前目录下。 # In[6]: from mindvision.engine.callback import ValAccMonitor import mindspore as ms num_epochs = 1 # 模型训练与验证,训练完成后保存验证精度最高的ckpt文件(best.ckpt)到当前目录下 model.train(num_epochs, dataset_train, callbacks=[ValAccMonitor(model, dataset_val, num_epochs), ms.TimeMonitor()]) # ### 可视化模型预测 # # 定义 `visualize_model` 函数,使用上述验证精度最高的模型对输入图像进行预测,并将预测结果可视化。 # In[8]: import matplotlib.pyplot as plt import numpy as np from PIL import Image import mindspore as ms def visualize_model(path): image = Image.open(path).convert("RGB") image = image.resize((224, 224)) plt.imshow(image) # 归一化处理 mean = np.array([0.485 * 255, 0.456 * 255, 0.406 * 255]) std = np.array([0.229 * 255, 0.224 * 255, 0.225 * 255]) image = np.array(image) image = (image - mean) / std image = image.astype(np.float32) # 图像通道由(h, w, c)转换为(c, h, w) image = np.transpose(image, (2, 0, 1)) # 扩展数据维数为(1, c, h, w) image = np.expand_dims(image, axis=0) # 定义并加载网络 net = mobilenet_v2(num_classes=6, resize=224) param_dict = ms.load_checkpoint("./best.ckpt") ms.load_param_into_net(net, param_dict) model = ms.Model(net) # 模型预测 pre = model.predict(ms.Tensor(image)) result = np.argmax(pre) #class_name = {0: "Croissants", 1: "Dog"} class_name = {0:"cardboard", 1:"glass", 2:"metal", 3:"paper", 4: "plastic", 5: "trash"} plt.title(f"Predict: {class_name[result]}") return result image1 = "Garbage_Data/infer/cardboard1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 1) visualize_model(image1) image2 = "Garbage_Data/infer/glass1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 2) visualize_model(image2) image3 = "Garbage_Data/infer/metal1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 2) visualize_model(image3) image4 = "Garbage_Data/infer/paper1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 2) visualize_model(image4) image5 = "Garbage_Data/infer/plastic1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 2) visualize_model(image5) image6 = "Garbage_Data/infer/trash1.jpg" plt.figure(figsize=(15, 7)) plt.subplot(1, 2, 2) visualize_model(image6) plt.show() # ### 模型导出 # # 在模型训练完后,训练完成后的网络模型(即CKPT文件)转换为MindIR格式,用于后续手机侧的推理。通过`export`接口会在当前目录下会生成`mobilenet_v2_1.0_224.mindir`文件。 # In[9]: import mindspore as ms # 定义并加载网络参数 net = mobilenet_v2(num_classes=6, resize=224) param_dict = ms.load_checkpoint("best.ckpt") ms.load_param_into_net(net, param_dict) # 将模型由ckpt格式导出为MINDIR格式 input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32) ms.export(net, ms.Tensor(input_np), file_name="mobilenet_v2_1.0_224", file_format="MINDIR")