基于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()