基于pytorch1.2版本的，个人数据读取，与tensorboard显示

1.定义自己的数据读取类：

class ListDataset(Dataset):
    def __init__(self, listpath, img_size=128):
        img_lists = []
        label_lists = []
        for sub1 in os.listdir(listpath):
            sub1_path = os.path.join(listpath, sub1)
            if not os.path.isdir(sub1_path):
                continue
            for sub2 in os.listdir(sub1_path):
                if not sub2.endswith('.png'):
                    continue
                img_path = os.path.join(sub1_path, sub2)
                img_lists.append(img_path)
                if sub1 == 'actors':
                    label_lists.append(0)
                elif sub1 == 'actresses':
                    label_lists.append(1)
                else:
                    raise ValueError("sub1 is wrong")

        self.img_files = img_lists
        self.label_lists = label_lists
        self.img_size = img_size

    def __getitem__(self, index):
        img_path = self.img_files[index % len(self.img_files)].rstrip()
        img = cv2.imread(img_path)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = transforms.ToTensor()(img)
        if len(img.shape) != 3:
            raise ValueError("img.shape is wrong")
        label = self.label_lists[index % len(self.img_files)]
        return img_path, img, label

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

说明： 1.读取图像位置，个人比较习惯使用opencv

             2.label位置，这个位置输出类型是个int，个人也有尝试改变其输出类型，最后发现这个位置和损失函数相关，如果使用预定于的交叉熵损失，int类型就可以，不同特意转换为numpy或者tensor，如果使用自定义损失函数的话，在于个人怎么改，loss的位置怎么接了。

             3.我看pytorch的yolov3程序，他在data读取环节定义了一个collate_fn函数，只在__getitem__函数中将数据产生list，然后在collate_fn中读取list中label的具体内容，然后在torch.utils.data.DataLoader(collate_fn=dataset.collate_fn）中将collate_fn付过去，也可以正常读取数据，个人稍微试了一下，也是没成功。

2.tensorboard：

from torch.utils import tensorboard

_log_dir = r'D:\tTest\m\test1'
if not os.path.exists(_log_dir):
    os.mkdir(_log_dir)
tsw = tensorboard.SummaryWriter(log_dir=_log_dir)
esw = tensorboard.SummaryWriter(log_dir=os.path.join(_log_dir, 'eval'))

tsw.add_scalar('train/loss', loss, i)

esw.add_scalar('val/loss', val_loss, i)
esw.add_image('val/image', imgs[0], i)

说明：1.pytorch1.2版本之后对tensorboard的支持真的是好优秀，不要再信那些低版本乱七八糟的tensorboardx什么之类的

            2.可简单了，就像上面的代码，图像的位置需要转换通道，pytorch图像的通道的CWH，显示出来的图像颜色不对。

3. 自定义model
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 8, kernel_size=3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(8, 4, kernel_size=3, stride=2, padding=1)
        self.fc1 = nn.Linear(4*8*8, 2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = x.view(x.size()[0], -1)
        out = self.fc1(x)
        return out

说明：这个没啥可说的，pytorch定义网络的时候，一些数需要自己计算好

4.完整代码：

from datasets import *
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import tensorboard
from torch.utils.data import Dataset
import os
import cv2
import torchvision.transforms as transforms

class ListDataset(Dataset):
    def __init__(self, listpath, img_size=128):
        img_lists = []
        label_lists = []
        for sub1 in os.listdir(listpath):
            sub1_path = os.path.join(listpath, sub1)
            if not os.path.isdir(sub1_path):
                continue
            for sub2 in os.listdir(sub1_path):
                if not sub2.endswith('.png'):
                    continue
                img_path = os.path.join(sub1_path, sub2)
                img_lists.append(img_path)
                if sub1 == 'actors':
                    label_lists.append(0)
                elif sub1 == 'actresses':
                    label_lists.append(1)
                else:
                    raise ValueError("sub1 is wrong")

        self.img_files = img_lists
        self.label_lists = label_lists
        self.img_size = img_size

    def __getitem__(self, index):
        img_path = self.img_files[index % len(self.img_files)].rstrip()
        img = cv2.imread(img_path)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = transforms.ToTensor()(img)
        if len(img.shape) != 3:
            raise ValueError("img.shape is wrong")
        label = self.label_lists[index % len(self.img_files)]
        return img_path, img, label

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 8, kernel_size=3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(8, 4, kernel_size=3, stride=2, padding=1)
        self.fc1 = nn.Linear(4*8*8, 2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = x.view(x.size()[0], -1)
        out = self.fc1(x)
        return out

def main():
    #data
    train_path = r'D:\tTest\data\train'
    train_dataset = ListDataset(train_path)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=8,
        shuffle=True,
    )
    valid_path = r'D:\tTest\data\valid'
    valid_dataset = ListDataset(valid_path)
    valid_dataloader = torch.utils.data.DataLoader(
        valid_dataset,
        batch_size=8,
        shuffle=True,
    )

    #tensorboard
    _log_dir = r'D:\tTest\m\test1'
    if not os.path.exists(_log_dir):
        os.mkdir(_log_dir)
    tsw = tensorboard.SummaryWriter(log_dir=_log_dir)
    esw = tensorboard.SummaryWriter(log_dir=os.path.join(_log_dir, 'eval'))

    #loss
    criterion = nn.CrossEntropyLoss()

    net = Net()

    optimizer = torch.optim.SGD(net.parameters(), lr=0.03)

    epochs = 100

    for i in range(epochs):
        loss_sum = 0
        for batch_i, (_, imgs, targets) in enumerate(train_dataloader):
            predict = net(imgs)
            loss = criterion(predict, targets)  
            loss_sum += loss
            if batch_i % 10 == 0:
                tsw.add_scalar('train/loss', loss, i)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(i, "train_loss_sum:", loss_sum.data)
        # valid
        if i % 3 == 0:
            loss_sum_val = 0.0
            for batch_i, (_, imgs, targets) in enumerate(valid_dataloader):
                predict = net(imgs)
                val_loss = criterion(predict, targets) 
                loss_sum_val += val_loss
                if batch_i % 10 == 0:
                    esw.add_scalar('val/loss', val_loss, i)
                    esw.add_image('val/image', imgs[0], i)
            print(i, "val_loss_sum:", loss_sum_val.data)

if __name__ == '__main__':
    main()