python制作图片数据集_Pytorch自己加载单通道图片用作数据集训练的实例

pytorch 在torchvision包里面有很多的的打包好的数据集，例如minist,Imagenet-12,CIFAR10 和CIFAR100。在torchvision的dataset包里面，用的时候直接调用就行了。具体的调用格式可以去看文档(目前好像只有英文的)。网上也有很多源代码。

不过，当我们想利用自己制作的数据集来训练网络模型时，就要有自己的方法了。pytorch在torchvision.dataset包里面封装过一个函数ImageFolder()。这个函数功能很强大，只要你直接将数据集路径保存为例如“train/1/1.jpg ,rain/1/2.jpg …… ”就可以根据根目录“./train”将数据集装载了。

?

但是后来我发现一个问题，就是这个函数加载出来的图像矩阵都是三通道的，并且没有什么参数调用可以让其变为单通道。如果我们要用到单通道数据集(灰度图)的话，比如自己加载Lenet-5模型的数据集，就只能自己写numpy数组再转为pytorch的Tensor()张量了。

接下来是我做的过程：

首先，还是要用到opencv,用灰度图打开一张图片，省事。

?

其次，pytorch有自己的numpy转Tensor函数，直接转就行了。

?

下一步利用torch.util和torchvision里面的dataLoader函数，就能直接得到和torchvision.dataset里面封装好的包相同的数据集样本了

?

最后就是自己建网络设计参数训练了，这部分和文档以及github中的差不多，就不赘述了。

下面是整个程序的源代码，我利用的还是上次的车标识别的数据集，一共分四类，用的是2层卷积核两层全连接。

源代码：

?

最后的训练结果和在keras下差不多，不过我训练的时候好像把训练集和测试集弄反了，数目好像测试集比训练集还多，有点尴尬，不过无伤大雅。结果图如下：

以上这篇Pytorch自己加载单通道图片用作数据集训练的实例就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持服务器之家。

原文链接：https://blog.csdn.net/qq_31548387/article/details/73302319