Dataset可以遍历数据集,每次输出一组(feature,label)。dataloader相当于是dataset的接口,顺便可以Dataset做一些调整,比如shuffle、batchsize……
因此,自定义dalaloader的关键是定义Dataset!
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
import pandas as pd
from torchvision.transforms import ToTensor, Lambda
from torchvision import datasets, models, transforms
class TrainDataset(Dataset):
def __init__(self, root_dir, csv_file, transform=None, target_transform=None):
self.root_dir = root_dir
self.labels = pd.read_csv(csv_file)
# features的归一化:
self.transform = transform
# labels的归一化:
self.target_transform = target_transform
def __len__(self):
return self.labels.shape[0]
def __getitem__(self, index):
img_name = os.path.join(self.root_dir,
self.labels.iloc[index, 0])
image = Image.open(img_name)
label = self.labels.iloc[index, 1:].astype(int).to_numpy()
label = list(label)[0].tolist()
if self.transform:
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)
# print(label)
return image, label
我这里root_dir=’'是因为我的label.csv文件中存储的是features的全路径:
# 1. Load and normalize training data
dataset = TrainDataset(
root_dir='',
csv_file='labels.csv',
transform=transforms.Compose([
transforms.Resize([224, 224]),
ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]),
target_transform=None
)
三通道彩色图像:
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
一通道的黑白图像:
transforms.Normalize((0.5,), (0.5,))
batch_size = 16
trainloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=1)
classes = ('s0', 's1', 's2', 's3', 's4', 's5', 's6', 's7', 's8', 's9')
我这个和pytorch CIFAR-10分类器教程设置的一样。
区别是self.fc1 = nn.Linear(16 * 53 * 53, 120),这里是165353,原本是1655。修改原因是本来samples的size为3232,我这里resize到224224,网络设置也要随之改变。
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 53 * 53, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 53 * 53)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
logits = x
return logits
# 2. Define a Convolutional Neural Network
net = Net()
# print(net)
# 3. Define a Loss function and optimizer
criterion = nn.CrossEntropyLoss()
# optimizer = optim.SGD(net.parameters(), lr=0.0001, momentum=0.9)
optimizer = optim.Adam(net.parameters(), lr=0.001)
神经网络+训练features+训练labels 都要放到GPU
# training on GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using {} device'.format(device))
net = net.to(device)
print('sent net to gpu')
# 4. Train the network
epoch1 = 10
for epoch in range(epoch1): # loop over the dataset multiple times
running_loss = 0.0
Loss = []
for i, (inputs, labels) in enumerate(trainloader, 0):
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
inputs = inputs.to(device)
labels = labels.to(device)
# outputs是网络的输出,用Softmax转变为该sample属于各类别的概率:
outputs = net(inputs)
loss = criterion(outputs, labels)
softmax = nn.Softmax(dim=1)
pred_probab = softmax(outputs)
nb_batches = pred_probab.shape[0]
nb_classes = pred_probab.shape[1]
max_idxes = []
for j in range(nb_batches):
max = pred_probab[j][0]
max_idx = 0
for k in range(nb_classes):
prob = pred_probab[j][k]
if prob > max:
max = prob
max_idx = k
max_idxes.append(max_idx)
print(labels)
pred_label = []
for ii in max_idxes:
pred_label.append(classes[ii])
print(pred_label)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
Loss.append(running_loss)
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
Loss0 = torch.tensor(Loss)
np.save('./loss/epoch_{}'.format(epoch), Loss0)
print('Finished Training')
确保网络能够在小训练集上完全拟合,验证有没有错误。
在这里我设计了一个小训练集,十类数据,每一类20个,所以一共20*10=200个数据。设置epoch=10,通过输出predicts和ground truth对比,可以看到完全拟合: