Pytorch使用本地图片实现mnist

一、下载图片解压即可

链接：https://pan.baidu.com/s/146E7BKbQzHO1q0r5x35Dgg 
提取码：bczs

二、代码实现

原始代码路径：https://github.com/pytorch/examples/tree/master/mnist，只要稍作修改即可，因为图片是RGB的，所以有2种方案

方案 1、直接使用RGB(修改如下代码)

1. 将channel数修改为3

self.conv1 = nn.Conv2d(3, 32, 3, 1)

2. 修改训练和测试代码

    train_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder('d:/mnist/train',
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder('d:/mnist/test', transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

 

方案2、RGB转换成灰阶图(修改如下代码)

    train_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder('d:/mnist/train',
                       transform=transforms.Compose([
                           transforms.Grayscale(num_output_channels=1),
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,)),
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder('d:/mnist/test', transform=transforms.Compose([
                           transforms.Grayscale(num_output_channels=1),
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)