写这个程序主要是为了熟悉pytorch自定义dataset的方式，所以没有用torchvision自带的cifar10，而是直接将原始图片都下下来了（下载：https://pjreddie.com/media/files/cifar.tgz）。
dataset实现方式也很简单，我个人喜欢将数据集的所有图像地址和所有图像标签分别存在两个npy文件中，然后两个numpy array作为dataset输入即可。
完整的程序实现起来很简单，一个py文件即可：
# -*- coding: utf-8 -*-
import os, torch, glob, time, copy, torchvision, operator, itertools
import numpy as np
from torch.autograd import Variable
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import torch.optim as optim
from torch.optim import lr_scheduler
from sklearn.cluster import KMeans
def get_label(str):
if 'airplane' in str:
return 0
elif 'automobile' in str:
return 1
elif 'bird' in str:
return 2
elif 'cat' in str:
return 3
elif 'deer' in str:
return 4
elif 'dog' in str:
return 5
elif 'frog' in str:
return 6
elif 'horse' in str:
return 7
elif 'ship' in str:
return 8
elif 'truck' in str:
return 9
class myDataset(Dataset):
# TensorDataset继承Dataset, 重载了__init__, __getitem__, __len__
# 实现将一组Tensor数据对封装成Tensor数据集
# 能够通过index得到数据集的数据，能够通过len，得到数据集大小
def __init__(self, files_list, labels_list, transform):
self.files_list = files_list
self.labels_list = labels_list
self.transform = transform
def __getitem__(self, index):
img = Image.open(self.files_list[index])
img = self.transform(img)
return img, self.labels_list[index]
def __len__(self):
return len(self.labels_list)
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
since2 = time.time()
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'val']:
if phase == 'train':
scheduler.step()
model.train(True) # Set model to training mode
else:
model.train(False) # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for step, (inputs, labels) in enumerate(dataloaders[phase]):
# get the inputs
# wrap them in Variable
if use_gpu:
inputs = Variable(inputs.cuda())
labels = Variable(labels.cuda())
else:
inputs, labels = Variable(inputs), Variable(labels)
# zero the parameter gradients
optimizer.zero_grad()
# forward
outputs = model(inputs)
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
if step % 100 ==0:
print(step, loss.data[0])
# statistics
running_loss += loss.data[0] * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects / dataset_sizes[phase]
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
torch.save(best_model_wts, 'model_ft.pth')
print(time.time() - since2)
if __name__ == '__main__':
# training_set_dir = './cifar/train'
# test_set_dir = './cifar/test'
#
# train_images_list = glob.glob(os.path.join(training_set_dir, '*.' + 'png'))
# test_images_list = glob.glob(os.path.join(test_set_dir, '*.' + 'png'))
#
# train_images_labels_list = [get_label(x) for x in train_images_list]
# test_images_labels_list = [get_label(x) for x in test_images_list]
#
# np.save('train_images_list.npy', train_images_list)
# np.save('test_images_list.npy', test_images_list)
#
# np.save('train_images_labels_list.npy', train_images_labels_list)
# np.save('test_images_labels_list.npy', test_images_labels_list)
train_images_list = np.load('train_images_list.npy')
test_images_list = np.load('test_images_list.npy')
train_images_labels_list = np.load('train_images_labels_list.npy')
train_images_labels_list = train_images_labels_list.astype('int64')
test_images_labels_list = np.load('test_images_labels_list.npy')
test_images_labels_list = test_images_labels_list.astype('int64')
num_classes = 10
BATCH_SIZE = 128
num_workers = 2
use_gpu = torch.cuda.is_available()
transform = transforms.Compose([
transforms.Resize((224, 224)),
# transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
]
)
train_dataset = myDataset(train_images_list, train_images_labels_list, transform)
test_dataset = myDataset(test_images_list, test_images_labels_list, transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE,
shuffle=True, num_workers=num_workers)
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE,
shuffle=False, num_workers=num_workers)
dataloaders = {}
dataloaders['train'] = train_loader
dataloaders['val'] = test_loader
dataset_sizes = {}
dataset_sizes['train'] = len(train_images_labels_list)
dataset_sizes['val'] = len(test_images_labels_list)
model_ft = models.resnet18(pretrained=False)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, num_classes)
if use_gpu:
model_ft = nn.DataParallel(model_ft)
model_ft = model_ft.cuda()
criterion = nn.CrossEntropyLoss()
# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.1)
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
num_epochs=2)
