pythorch CNN 小教程，关于如何自定义数据集

之前用官方的minist数据集测试过了，现在想加载自己的数据集并进行识别分类，应该如何加载呢？

利用ImageFolder函数，将所需分类的图片提前按文件夹分类分别置于train val文件夹下即可

代码如下所示：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy
import cv2
import torchvision
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.autograd import Variable

class CNNnet(torch.nn.Module):
    def __init__(self):
        super(CNNnet, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=3,    # input 3,16,400,1400
                            out_channels=16, #16个filter，提取16个特征作为输出
                            kernel_size=5,
                            stride=1,
                            padding=2),
            torch.nn.BatchNorm2d(16),  # 是一个filter,提取特征， 特征窗
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)  # 池化层 out 16,200,700
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(16, 32, 5, 1, 2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)  # 32,100,350
        )
        self.conv3 = torch.nn.Sequential(
            torch.nn.Conv2d(32, 32, 5, 1, 2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)  # 32 50 175
        )
        # self.conv4 = torch.nn.Sequential(
        #
        #     torch.nn.Conv2d(32, 64, 5, 1, 2),  # 2*32*48
        #     torch.nn.BatchNorm2d(64),
        #     torch.nn.ReLU(),
        #     torch.nn.MaxPool2d(2)  # 64,25,87
        # )
        # self.mlp1 = torch.nn.Linear(2*2*2*16*24, 2000)
        self.mlp1 = torch.nn.Linear(32*64*64, 1000)
        self.mlp2 = torch.nn.Linear(1000, 3)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)

        # x = self.conv4(x)
        # print(x.shape)
        # t = x.size(0)
        x = self.mlp1(x.view(x.size(0), -1))
        out = self.mlp2(x)
        return out


train_data = torchvision.datasets.ImageFolder('D:/pytorchdemo/pic/train',
                                            transform=transforms.Compose([
                                                # transforms.Scale(512,768),
                                                # transforms.Resize((512,768)),
                                                transforms.CenterCrop(512),
                                                transforms.ToTensor()])
                                            )
test_data = torchvision.datasets.ImageFolder('D:/pytorchdemo/pic/val',
                                            transform=transforms.Compose([
                                                # transforms.Scale(512,768),
                                                # transforms.Resize((768,512)),
                                                transforms.CenterCrop(512),
                                                transforms.ToTensor()])
                                            )
train_loader = torch.utils.data.DataLoader(train_data, batch_size=2,shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data,batch_size=2, shuffle=True)
# print(train_data.train_data.size())
x1,x2 = train_data.samples[0]
# print(x1)


model = CNNnet()
loss_func = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)  # 设置学习方法
# num_epochs = 100 # 设置训练次数
EPOCH = 20
# print(model)

def train_fun():
    # epochs = 0

    loss_list = []
    for epoch in range(EPOCH):
        step = 0
        for data in (train_loader):
            b_x,b_y=data
            b_x,b_y=Variable(b_x),Variable(b_y)
            out_put=model(b_x)
            loss=loss_func(out_put,b_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            step += 1
        #
        # if epochs < 20:
        #     epochs += 1
            print('Epoch: ', epoch,'step: ',step,'loss: ', float(loss))
            loss_list.append(float(loss))
        # else:
        #     epochs += 1

    return loss_list


# criterion = nn.MSELoss()  # 设定误差函数
def test_fun():
    eval_loss = 0
    eval_acc = 0
    # for eopch in range(EPOCH):
    for data in (test_loader):
        b_x, b_y = data
        b_x, b_y = Variable(b_x), Variable(b_y)  # b_x 是 img b_y是标签
        out_put = model(b_x)
        loss = loss_func(out_put, b_y)
        eval_loss += loss.data.item() * b_y.size(0)
        _, pred = torch.max(out_put, 1)
        # print("pred",pred,b_y)
        num_correct = (pred == b_y).sum()
        eval_acc += num_correct.item()
        # print(eval_acc)
        print('Test Loss: {:.4f}, Acc: {:.4f}'.format(
        eval_loss / (len(test_data)),
        eval_acc / (len(test_data)), 100. * eval_acc / len(test_data)
    ))


def main():
    print("------starting-------")


if __name__ == '__main__':
    main()
    print("ok")
    # params = list(model.parameters())
    train_fun()
    #torch.save(model,"mycnn2")
    # test_fun()