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

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

今天小编就为大家分享一篇Pytorch自己加载单通道图片用作数据集训练的实例，具有很好的参考价值，希望对

大家有所帮助。一起跟随小编过来看看吧

pytorch 在torchvision包里面有很多的的打包好的数据集，例如minist,Imagenet-12,CIFAR10 和CIFAR100。在torchvision的

dataset包里面，用的时候直接调用就行了。具体的调用格式可以去看文档（目前好像只有英文的）。网上也有很多源代码。

不过，当我们想利用自己制作的数据集来训练网络模型时，就要有自己的方法了。pytorch在torchvision.dataset包里面封装过

一个函数ImageFolder（）。这个函数功能很强大，只要你直接将数据集路径保存为例如“train/1/1.jpg ,rain/1/2.jpg …… ”就可

以根据根目录“./train”将数据集装载了。

#读取图片　这里是灰度图　

for item in all_path:

img = cv2.imread(item[1],0)

img = cv2.resize(img,(28,28))

arr = np.asarray(img,dtype="float32")

data_x[i ,:,:,:] = arr

i+=1

data_y.append(int(item[0]))

data_x = data_x / 255

dataset = dataf.TensorDataset(data_x,data_y)

loader = dataf.DataLoader(dataset, batch_size=batchsize, shuffle=True)

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

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

源代码：

# coding=utf-8

import os

import cv2

import numpy as np

import random

import torch.optim as optim

#训练参数

cuda = False

train_epoch = 20

train_lr = 0.01

train_momentum = 0.5

batchsize = 5