Pytorch_加载自己的数据集并训练

如何制作数据集txt,并进行加载

1. 根据图片制作train.txt val.txt 和 test.txt

# MakeTxt.py
"""
原始图片保存路径
**************文件保存方式  一  *******************
--data
    --train
        --category_1
        --category_2
        ...
    --val
        --category_1
        --category_2
        ...
    --test
        --category_1
        --category_2
        ...

**************文件保存方式  二  *******************
--data
    --category_1
        --001.jpg
        --002.jpg
        ...
    --category_2
        --001.jpg
        --002.jpg
        ...
    ...

"""

import os
# 【1. 获取文件绝对路径】
def get_files_path(file_dir):
    dir_names = []
    for fn in os.listdir(file_dir): #fn 表示的是文件名
        dir_names.append(os.path.join(file_dir,fn))
    print(dir_names)
    return dir_names
def get_dir_img(file_dir):
    filenames = []
    for root, dirs, files in os.walk(file_dir):
        for name in files:
            # print(os.path.join(root, name))
            filenames.append(os.path.join(root, name))
    return filenames
def make_txt(img_root,txt_root,quantity_proportion):
    """
    img_root : 图像保存路径
    txt_root : txt路径
    quantity_proportion : [train,eval,test] 各数据集的比例 eg. [0.8, 0.1, 0.1]
    """
    # 创建 txt文件
    txt_name = [txt_root + '/train.txt', txt_root +'/val.txt', txt_root + '/test.txt']
    for file_path in txt_name:
        if os.path.isfile(file_path):
            os.remove(file_path)
    train = open(txt_name[0],'a')
    val = open(txt_name[1], 'a')
    test = open(txt_name[2], 'a')
    sort_files = get_files_path(img_root)
    for i,file in enumerate(sort_files):	# 进行枚举，i为索引
        tem_total_img = get_dir_img(file)
        num_img = len(tem_total_img)
        span_num = [int(x* num_img) for x in quantity_proportion] 
        print(span_num,num_img)
        for j in range(span_num[0]+1):
            train.write(tem_total_img[j] + ' ' + str(i) + '\n')
        for k in range(span_num[0]+1,span_num[0]+span_num[1]+1):
            val.write(tem_total_img[k] + ' ' + str(i) + '\n')
        for p in range(span_num[0]+span_num[1]+1,num_img):
            test.write(tem_total_img[p] + ' ' + str(i) + '\n')
        train.close()
        val.close()
        test.close()

def main():
    quantity_proportion = [0.8, 0.1, 0.1]
    make_txt('./pokeman','./dataTxt',quantity_proportion)
    pass
if __name__=='__main__':
    main()

2. torch.utils.data.Dataset用来制作数据集

Image.open(path).convert(‘RGB’) 图片转换成RGB通道，对应后面的网络输入channel改为3

#***************************一些必要的包的调用********************************
import torch.nn.functional as F
import torch
import torch 
import torch.nn as nn
from torch.autograd import Variable
import torchvision.models as models
from torchvision import transforms, utils
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
import torch.optim as optim
import os
#***************************初始化一些函数********************************
gpu_id =0
torch.cuda.set_device(gpu_id)#使用GPU
learning_rate = 0.0001       #学习率的设置

#*************************************数据集的设置****************************************************************************
    #定义读取文件的格式
def default_loader(path):
    return Image.open(path).convert('RGB')

class MyDataset(Dataset): 
    #创建自己的类： MyDataset,这个类是继承的torch.utils.data.Dataset
    #**********************************  #使用__init__()初始化一些需要传入的参数及数据集的调用**********************
	def __init__(self, txt, transform=None,target_transform=None, loader=default_loader):	                          
	    # super(MyDataset, self).__init__()
		#对继承自父类的属性进行初始化
		fh = open(txt, 'r') 
		imgs = []  
		# 按照传入的路径和txt文本参数以只读的方式打开这个文本
		for line in fh: #迭代该列表#按行循环txt文本中的内
			line = line.strip('\n')
			line = line.rstrip('\n')# 删除 本行string 字符串末尾的指定字符，这个方法的详细介绍自己查询python
			words = line.split() #用split将该行分割成列表  split的默认参数是空格，所以不传递任何参数时分割空格
			imgs.append((words[0],int(words[1])))
            #把txt里的内容读入imgs列表保存，具体是words几要看txt内容而定 
            # 很显然，根据我刚才截图所示txt的内容，words[0]是图片信息，words[1]是lable       
		self.imgs = imgs
		self.transform = transform
		self.target_transform = target_transform
		self.loader = loader        
    #*************************** #使用__getitem__()对数据进行预处理并返回想要的信息**********************
	def __getitem__(self, index):#这个方法是必须要有的，用于按照索引读取每个元素的具体内容
		fn, label = self.imgs[index]
		#fn是图片path #fn和label分别获得imgs[index]也即是刚才每行中word[0]和word[1]的信息
		img = self.loader(fn) # 按照路径读取图片
		if self.transform is not None:
			img = self.transform(img) #数据标签转换为Tensor
		return img,label
	#return回哪些内容，那么我们在训练时循环读取每个batch时，就能获得哪些内容
	#**********************************  #使用__len__()初始化一些需要传入的参数及数据集的调用**********************
	def __len__(self):
	        #这个函数也必须要写，它返回的是数据集的长度，也就是多少张图片，要和loader的长度作区分
		return len(self.imgs)


def my_data_set(txt_root = './dataTxt/',batchsize=32,resize=227):
	#图像的初始化操作

	_transforms = transforms.Compose([
		transforms.RandomResizedCrop((resize,resize)),
		transforms.ToTensor(),])
	train_data = MyDataset(txt=txt_root + 'train.txt', transform=_transforms)
	test_data  = MyDataset(txt=txt_root + 'test.txt', transform=_transforms)
	val_data   = MyDataset(txt=txt_root + 'val.txt', transform=_transforms)
	train_loader = DataLoader(dataset=train_data, batch_size=batchsize, shuffle=True,num_workers=1)
	test_loader  = DataLoader(dataset=test_data, batch_size=batchsize, shuffle=False,num_workers=1)
	val_loader   = DataLoader(dataset=val_data, batch_size=batchsize, shuffle=False,num_workers=1)
	print('num_of_trainData:', len(train_data))
	print('num_of_testData:', len(test_data))
	print('num_of_valData:', len(val_data))
	return train_loader,test_loader,val_loader
	# print(train_data)

搭建RestNet18进行测试

import torch
import time
from torch import nn,optim
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
import MyDataset as md
import sys
sys.path.append("..")
from Pytorch_learning import pytorch_deep as pyd
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 创建残差块 rest block
class Residual(nn.Module):
    def __init__(self, in_channels,out_channels,use_1x1conv = False,stride = 1):
        super(Residual,self).__init__()
        self.conv1 = nn.Conv2d(in_channels,out_channels,kernel_size=3,padding = 1, stride = stride)
        self.conv2 = nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(in_channels,out_channels,kernel_size=3,padding = 1, stride = stride)
        else:self.conv3 = None
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
    def forward(self,X):
        Y = F.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
        if self.conv3:
            X = self.conv3(X)
        return F.relu(X+Y)
def resnet_block(in_channels, out_channels, num_residuals,first_block=False):
    if first_block:
        assert in_channels == out_channels # 第⼀个模块的通道数同输⼊通道数⼀致
    blk = []
    for i in range(num_residuals):
        if i == 0 and not first_block:
            blk.append(Residual(in_channels, out_channels,use_1x1conv=True, stride=2))
        else:
            blk.append(Residual(out_channels, out_channels))
    return nn.Sequential(*blk)

def my_restnet_18(num_classes= 6):
    net = nn.Sequential(
                nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
                nn.BatchNorm2d(64),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=3, stride=2, padding=1))

    net.add_module("resnet_block1", resnet_block(64, 64, 2,first_block=True))
    net.add_module("resnet_block2", resnet_block(64, 128, 2))
    net.add_module("resnet_block3", resnet_block(128, 256, 2))
    net.add_module("resnet_block4", resnet_block(256, 512, 2))
    net.add_module("global_avg_pool", pyd.GlobalAvgPool2d()) #GlobalAvgPool2d的输出: (Batch, 512, 1, 1)
    net.add_module("fc", nn.Sequential(pyd.FlattenLayer(),nn.Linear(512, num_classes)))
    return net

def my_train(save_path = './weight/restNet18_3.pth',resize=96,batch_size = 32):
    ## Training
    net = my_restnet_18(num_classes= 6)
    # 如出现“out of memory”的报错信息，可减⼩batch_size或resize
    train_iter, test_iter,val_iter =md.my_data_set(batchsize=batch_size,resize=resize)
    # train_iter, test_iter = pyd.load_data_fashion_mnist(batch_size,resize)
    lr, num_epochs = 0.001, 10
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    ###训练过程
    net = net.to(device)
    print("training on ", device)
    ## 损失函数 交叉熵损失
    loss = torch.nn.CrossEntropyLoss()  
    batch_count = 0
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, start = 0.0, 0.0, 0,time.time()
        for X, y in train_iter:
            # print(X)
            # print(len(X.data))
            X = X.to(device)   # 数据放到GPU
            y = y.to(device)
            y_hat = net(X)     #得到网络输出结果
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_l_sum += l.cpu().item()
            train_acc_sum += (y_hat.argmax(dim=1) ==y).sum().cpu().item()
            n += y.shape[0]
            batch_count += 1
        test_acc = pyd.evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f,time %.1f sec'% 
              (epoch + 1, train_l_sum / batch_count,train_acc_sum / n, test_acc, time.time() - start))
    # torch.save(net,save_path)
def my_test(pth_path = './weight/restNet18.pth',resize = 96,batch_size = 32):
    model_net = torch.load(pth_path)
    train_iter, test_iter,val_iter =md.my_data_set(batchsize=batch_size,resize=resize)
    # 预测正确的数量和总数量
    correct = 0
    total = 0
    # 使用torch.no_grad的话在前向传播中不记录梯度，节省内存
    # cv2.namedWindow('predictPic', cv2.WINDOW_NORMAL)
    # to_pil_image  = transforms.ToPILImage()
    with torch.no_grad():
        for images,labels in val_iter:
            # images, labels = data
            # print(images)
            print(len(images.data))
            images, labels = images.to(device), labels.to(device)
            # 预测
            # outputs = self.net(images)
            outputs = model_net(images)
            # 我们的网络输出的实际上是个概率分布，去最大概率的哪一项作为预测分类
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            # print(images.data[0])
            # print(len(images.data[0]))
            for i in range(len(images.data)):
                print('实际标签 {}，预测标签 {}'.format(labels[i],predicted[i]))

def test_one_img(img_path, pth_path = './weight/restNet18_2.pth',resize = 96):
    img = Image.open(img_path).convert('RGBA')
    test_transforms = transforms.Compose([
            transforms.RandomResizedCrop((resize,resize)),
            transforms.ToTensor(),])
    img = test_transforms(img) 
    img = torch.unsqueeze(img, 0) 
    model_net = torch.load(pth_path)

    with torch.no_grad():
            img = img.to(device)
            # 预测
            # outputs = self.net(images)
            outputs = model_net(img)
            print(outputs)
            # # 我们的网络输出的实际上是个概率分布，去最大概率的哪一项作为预测分类
            _, predicted = torch.max(outputs.data, 1)
            print('predicted = {}'.format(predicted.item()))
            # total += labels.size(0)
            # correct += (predicted == labels).sum().item()

def main():
    save_path = './weight/restNet18_3.pth'
    pth_path = save_path
    resize = 227
    batch_size = 32
    #训练
    my_train(save_path,resize,batch_size)
    # 测试
    # my_test(pth_path,resize,batch_size)
    # 测试一张图片
    # img_path = './pokeman/mewtwo/00000036.jpg'
    # test_one_img(img_path,pth_path,resize)
    pass
if __name__=='__main__':
    main()

相关调用

# -*- coding: utf-8 -*-
import torch
from torch import nn
import matplotlib.pyplot as plt
from IPython import display
import torchvision
import torchvision.transforms as transforms
import time
import torch.nn.functional as F
def use_svg_display():
    # ⽤⽮量图显示
    display.set_matplotlib_formats('svg')
def set_figsize(figsize=(3.5, 2.5)):
    use_svg_display()
    # 设置图的尺⼨
    plt.rcParams['figure.figsize'] = figsize

#小批量数据读取
def data_iter(batch_size, features, labels):
    num_examples = len(features)
    indices = list(range(num_examples))
    random.shuffle(indices)
    for i in range(0,num_examples,batch_size):
#         print(i)
        j = torch.LongTensor(indices[i:min(i +batch_size,num_examples)])
        yield features.index_select(0,j),labels.index_select(0,j)
# 定义损失函数
def squared_loss(y_hat, y):
    return (y_hat - y.view(y_hat.size())) ** 2 /2
# 优化函数
def sgd(params, lr, batch_size):
    for param in params:
        param.data -= lr * param.grad / batch_size

#       MNIST_FASHION 标签
def get_fashion_mnist_labels(labels):
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress','coat','sandal', 'shirt', 'sneaker', 'bag', 'ankleboot']
    label_list = []
    for item in labels:
        label_list.append(text_labels[int(item)])
#     return [text_labels[int(i)] for i in labels]
    return label_list


def show_fashion_mnist(images, labels):
    use_svg_display()
    # 这⾥的_表示我们忽略（不使⽤）的变量
    _, figs = plt.subplots(1, len(images), figsize=(12, 12))
    for f, img, lbl in zip(figs, images, labels):
        f.imshow(img.view((28, 28)).numpy())
        f.set_title(lbl)
        f.axes.get_xaxis().set_visible(False)
        f.axes.get_yaxis().set_visible(False)
    plt.show()

def load_data_fashion_mnist(batch_size = 256,resize=None,num_workers = 0):
    trans = []
    if resize:
        trans.append(torchvision.transforms.Resize(size=resize))
    trans.append(torchvision.transforms.ToTensor())
    transform = torchvision.transforms.Compose(trans)
    pwd_MNIST = '/home/tjmt/tjmt_new/notebook/Wangjiale_DeepLab/Pytorch_learning'
    mnist_train = torchvision.datasets.FashionMNIST(root=pwd_MNIST + '/MNIST', train=True, download=True,
                                                    transform=transform)
    mnist_test = torchvision.datasets.FashionMNIST(root=pwd_MNIST + '/MNIST', train=False, download=True,
                                                   transform=transform)
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return train_iter,test_iter
def evaluate_accuracy(data_iter,net):
#     模型net 在数据集data_iter 上的分类准确率
    acc_sum,n = 0.0,0
    for X,y in data_iter:
#         print((net(X).argmax(dim=1)==y).float().sum().item())
        acc_sum += (net(X).argmax(dim=1)==y).float().sum().item()
        n+=y.shape[0]
    return acc_sum/n
def train_ch3(net, train_iter,test_iter,loss,num_epochs,batch_size,params=None,lr=None,optimizer = None):
    for epoch in range(num_epochs):
        #模型训练次数 5次
        train_l_num, train_acc_num,n = 0.0,0.0,0
        for X,y in train_iter:
            #X 为小批量256个图像 28*28 y为标签  
            # 计算X softmax下的值   与损失函数值
            y_hat = net(X) 
            l = loss(y_hat,y).sum()
            
            #梯度清零
            if optimizer is not None:
                optimizer.zero_grad()
            elif params is not None and params[0].grad is not None:
                for param in params:
                    param.grad.data.zero_()
            l.backward()
            if optimizer is None:
                sgd(params,lr,batch_size)
            else:
                optimizer.step()
            train_l_num += l.item()
            train_acc_num += (y_hat.argmax(dim=1)==y).sum().item()
            n+= y.shape[0]
        test_acc = evaluate_accuracy(test_iter,net)
        print('epoch %d, loss %.4f,train_acc %.3f,test_acc %.3f'%(epoch+1,train_l_num/n, train_acc_num/n, test_acc))
class FlattenLayer(nn.Module):
    def __init__(self):
        super(FlattenLayer, self).__init__()
    def forward(self, x): # x shape: (batch, *, *, ...)
        return x.view(x.shape[0], -1)
def fun1():
    print('sssss')
    
# 计算训练准确度 20200625
def evaluate_accuracy(data_iter, net, device= torch.device('cuda' if torch.cuda.is_available() else 'cpu')):
    acc_sum,n = 0.0,0.0       #acc_sum 计算训练过程中模型输出结果与标签相同的数量
    with torch.no_grad():     # torch.no_grad() 该作用域下 的计算不会被 track
        for X, y in data_iter:
            if isinstance(net, torch.nn.Module):
                net.eval() # 评估模式, 这会关闭dropout
                acc_sum += (net(X.to(device)).argmax(dim=1) ==y.to(device)).float().sum().cpu().item()
                net.train() # 改回训练模式
            else:print('error net')
            n += y.shape[0]
    return acc_sum / n
def train_ch5(net, train_iter, test_iter, batch_size, optimizer, device, num_epochs):
    net = net.to(device)
    print("training on 111000", device)
    ## 损失函数 交叉熵损失
    loss = torch.nn.CrossEntropyLoss()  
    batch_count = 0
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, start = 0.0, 0.0, 0,time.time()
        for X, y in train_iter:
            X = X.to(device)   # 数据放到GPU
            y = y.to(device)
            y_hat = net(X)     #得到网络输出结果
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_l_sum += l.cpu().item()
            train_acc_sum += (y_hat.argmax(dim=1) ==y).sum().cpu().item()
            n += y.shape[0]
            batch_count += 1
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f,time %.1f sec'% 
              (epoch + 1, train_l_sum / batch_count,train_acc_sum / n, test_acc, time.time() - start))
    torch.save(net,'./mnist_alexnet_model.pth')
    # torch.save(net.state_dict(),'./mnist_alexnet.pth')
class GlobalAvgPool2d(nn.Module):
    # 全局平均池化层可通过将池化窗⼝形状设置成输⼊的⾼和宽实现
    def __init__(self):
        super(GlobalAvgPool2d, self).__init__()
    def forward(self, x):
        return F.avg_pool2d(x, kernel_size=x.size()[2:])
if __name__ == '__main__':
    main()