pytorch + visdom CNN处理自建图片数据集

最新推荐文章于 2024-08-24 11:47:11 发布

泛泛之素

最新推荐文章于 2024-08-24 11:47:11 发布

阅读量7.2k

点赞数 3

分类专栏：可视化机器学习 pytorch+visdom

本文链接：https://blog.csdn.net/tonydz0523/article/details/79075077

版权

该博客介绍了在Windows环境下，使用PyTorch和Visdom进行自定义图像数据集（蚂蚁和蜜蜂分类）的CNN模型训练。通过ImageFolder模块导入和预处理数据，创建了修改过的CNN网络。经过10和20个epoch的训练，准确率分别达到74.5%。进一步使用预训练的resnet18模型，10个epoch后准确率达到95%。

摘要由CSDN通过智能技术生成

环境

系统：win10
cpu：i7-6700HQ
gpu：gtx965m
python : 3.6
pytorch ：0.3

数据下载

来源自Sasank Chilamkurthy 的教程；数据：下载链接。
下载后解压放到项目根目录：
这里写图片描述
数据集为用来分类蚂蚁和蜜蜂。有大约120个训练图像，每个类有75个验证图像。

数据导入

可以使用 torchvision.datasets.ImageFolder(root,transforms) 模块可以将图片转换为 tensor。
先定义transform：

data_transforms = {
    'train': transforms.Compose([
        # 随机切成224x224 大小图片 统一图片格式
        transforms.RandomResizedCrop(224),
        # 图像翻转
        transforms.RandomHorizontalFlip(),
        # totensor 归一化(0,255) >> (0,1)   normalize   channel=（channel-mean）/std
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ]),
    "val" : transforms.Compose([
        # 图片大小缩放 统一图片格式
        transforms.Resize(256),
        # 以中心裁剪
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
}

导入，加载数据：

data_dir = './hymenoptera_data'
# trans data
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']}
# load data
data_loaders = {x: DataLoader(image_datasets[x], batch_size=BATCH_SIZE, shuf