Pytorch 使用CNN图像分类的实现

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
from PIL import Image

import csv
import collections
import os
import shutil

def buildDataset(root,dataType,dataSize):
  """构造数据集
  构造的图片存到root/{dataType}Data
  图片地址和标签的csv文件存到 root/{dataType}DataInfo.csv
  Args:
    root:str
      项目目录
    dataType:str
      'train'或者‘test'
    dataNum:int
      数据大小
  Returns:
  """
  dataInfo = []
  dataPath = f'{root}/{dataType}Data'
  if not os.path.exists(dataPath):
    os.makedirs(dataPath)
  else:
    shutil.rmtree(dataPath)
    os.mkdir(dataPath)
    
  for i in range(dataSize):
    # 创建0，1 数组
    imageArray=np.random.randint(0,2,(4,4))
    # 计算0，1数量得到标签
    allBlackNum = collections.Counter(imageArray.flatten())[0]
    innerBlackNum = collections.Counter(imageArray[1:3,1:3].flatten())[0]
    label = 0 if (allBlackNum-innerBlackNum)>innerBlackNum else 1
    # 将图片保存
    path = f'{dataPath}/{i}.jpg'
    dataInfo.append([path,label])
    im = Image.fromarray(np.uint8(imageArray*255))
    im = im.convert('1') 
    im.save(path)
  # 将图片地址和标签存入csv文件
  filePath = f'{root}/{dataType}DataInfo.csv'
  with open(filePath, 'w') as f:
    writer = csv.writer(f)
    writer.writerows(dataInfo)

root=r'/Users/null/Documents/PythonProject/Classifier'

buildDataset(root,'train',20000)

buildDataset(root,'test',10000)

class MyDataset(torch.utils.data.Dataset):

  def __init__(self, root, datacsv, transform=None):
    super(MyDataset, self).__init__()
    with open(f'{root}/{datacsv}', 'r') as f:
      imgs = []
      # 读取csv信息到imgs列表
      for path,label in map(lambda line:line.rstrip().split(','),f):
        imgs.append((path, int(<a style="color:transparent">本文来源gao($daima.com搞@代@#码(网5</a>label)))
    self.imgs = imgs
    self.transform = transform if transform is not None else lambda x:x
    
  def __getitem__(self, index):
    path, label = self.imgs[index]
    img = self.transform(Image.open(path).convert('1'))
    return img, label

  def __len__(self):
    return len(self.imgs)

trainData=MyDataset(root = root,datacsv='trainDataInfo.csv', transform=transforms.ToTensor())
testData=MyDataset(root = root,datacsv='testDataInfo.csv', transform=transforms.ToTensor())

import itertools

def chooseData(dataset,scale):
  # 将类别为1的排序到前面
  dataset.imgs.sort(key=lambda x:x[1],reverse=True)
  # 获取类别1的数目 ，取scale倍的数组，得数据不那么偏斜
  trueNum =collections.Counter(itertools.chain.from_iterable(dataset.imgs))[1]
  end = min(trueNum*scale,len(dataset))
  dataset.imgs=dataset.imgs[:end]
scale = 4
chooseData(trainData,scale)
chooseData(testData,scale)
len(trainData),len(testData)
(2250, 1122)

import torch.utils.data as Data

# 超参数
batchSize = 50
lr = 0.1
numEpochs = 20

trainIter = Data.DataLoader(dataset=trainData, batch_size=batchSize, shuffle=True)
testIter = Data.DataLoader(dataset=testData, batch_size=batchSize)