自己数据集制作
这里以kaggle比赛猫狗两类数据集为例，
首先根据数据集制作一个txt的路径文档，代码如下
import os
from PIL import Image
path='E:\dataset\kaggle_cat_dog\PetImages/'
filetrain=open("train.txt",'w')
filetest=open("test.txt",'w')
num=0
for class_pet in os.listdir(path):
for img in os.listdir(os.path.join(path,class_pet)):
_, ext = os.path.splitext(img)
absolute_path=os.path.join(path,class_pet,img)
if ext=='jpg'or 'png'and num%8!=0:
filetrain.writelines(absolute_path+" 0"+'\n') if class_pet=='Cat' else filetrain.writelines(absolute_path+" 1"+"\n")
elif ext=='jpg' or 'png' and num%8==0:
filetest.writelines(absolute_path + " 0" + '\n') if class_pet == 'Cat' else filetest.writelines(
absolute_path + " 1" + "\n")
num+=1
filetrain.close()
filetest.close()其中通过num%8将数据集按照8:1分成了训练集与测试集。最终得到了两个txt文件夹。
然后将数据集加载并读取，代码为
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import DataLoader,Dataset
class MyDataset(Dataset):
def __init__(self,root,datatxt,transform=None):
super(MyDataset,self).__init__()
fh=open(root+datatxt,'r')
imgs=[]
for line in fh:
line=line.rstrip()
words=line.split()
imgs.append((words[0],int(words[1])))
self.imgs=imgs
self.transform=transform
def __getitem__(self,index):
fn,label=self.imgs[index]
img=Image.open(fn).convert('RGB')
if self.transform is not None:
img=self.transform(img)
return img,label
def __len__(self):
return len(self.imgs)
transform_train = transforms.Compose([
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
trainset=MyDataset(root='E:/tensorflow\SENet-Tensorflow-master/',datatxt='train.txt',transform=transform_train)
trainloader=DataLoader(trainset,batch_size=8,shuffle=True)
testset=MyDataset(root='E:/tensorflow\SENet-Tensorflow-master/',datatxt='test.txt',transform=transform_test)
testloader=DataLoader(testset,batch_size=8,shuffle=True)接下来是添加网络模型，以vgg16为例
import torch
import torch.nn as tnn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
class VGG16(tnn.Module):
def __init__(self):
super(VGG16, self).__init__()
self.layer1 = tnn.Sequential(
# 1-1 conv layer
tnn.Conv3d(3, 64, kernel_size=3, padding=1),
tnn.BatchNorm3d(64),
tnn.ReLU(),
# 1-2 conv layer
tnn.Conv3d(64, 64, kernel_size=3, padding=1),
tnn.BatchNorm3d(64),
tnn.ReLU(),
# 1 Pooling layer
tnn.MaxPool3d(kernel_size=2, stride=2))
self.layer2 = tnn.Sequential(
# 2-1 conv layer
tnn.Conv3d(64, 128, kernel_size=3, padding=1),
tnn.BatchNorm3d(128),
tnn.ReLU(),
# 2-2 conv layer
tnn.Conv3d(128, 128, kernel_size=3, padding=1),
tnn.BatchNorm3d(128),
tnn.ReLU(),
# 2 Pooling lyaer
tnn.MaxPool3d(kernel_size=2, stride=2))
self.layer3 = tnn.Sequential(
# 3-1 conv layer
tnn.Conv3d(128, 256, kernel_size=3, padding=1),
tnn.BatchNorm3d(256),
tnn.ReLU(),
# 3-2 conv layer
tnn.Conv3d(256, 256, kernel_size=3, padding=1),
tnn.BatchNorm3d(256),
tnn.ReLU(),
# 3 Pooling layer
tnn.MaxPool3d(kernel_size=2, stride=2))
self.layer4 = tnn.Sequential(
# 4-1 conv layer
tnn.Conv3d(256, 512, kernel_size=3, padding=1),
tnn.BatchNorm3d(512),
tnn.ReLU(),
# 4-2 conv layer
tnn.Conv3d(512, 512, kernel_size=3, padding=1),
tnn.BatchNorm3d(512),
tnn.ReLU(),
# 4 Pooling layer
tnn.MaxPool3d(kernel_size=2, stride=2))
self.layer5 = tnn.Sequential(
# 5-1 conv layer
tnn.Conv3d(512, 512, kernel_size=3, padding=1),
tnn.BatchNorm3d(512),
tnn.ReLU(),
# 5-2 conv layer
tnn.Conv3d(512, 512, kernel_size=3, padding=1),
tnn.BatchNorm3d(512),
tnn.ReLU(),
# 5 Pooling layer
tnn.MaxPool3d(kernel_size=2, stride=2))
self.layer6 = tnn.Sequential(
# 6 Fully connected layer
# Dropout layer omitted since batch normalization is used.
tnn.Linear(4096, 4096),
tnn.BatchNorm1d(4096),
tnn.ReLU())
self.layer7 = tnn.Sequential(
# 7 Fully connected layer
# Dropout layer omitted since batch normalization is used.
tnn.Linear(4096, 4096),
tnn.BatchNorm1d(4096),
tnn.ReLU())
self.layer8 = tnn.Sequential(
# 8 output layer
tnn.Linear(4096, 2),
tnn.BatchNorm1d(2),
tnn.Softmax())
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = self.layer5(out)
vgg16_features = out.view(out.size(0), -1)
out = self.layer6(vgg16_features)
out = self.layer7(out)
out = self.layer8(out)
return vgg16_features, out模型定义结束后，开始训练：
import time
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def accuracy(output, target, topk=(1,)):
"""Computes the precision@k for the specified values of k"""
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def train(train_loader, model, criterion, optimizer, epoch):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(trainloader):
target = target.cuda(async=True)
input = input.cuda()
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
# compute output
output = model(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1 = accuracy(output.data, target, topk=(1,))[0]
losses.update(loss.data[0], input.size(0))
top1.update(prec1[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
epoch, i, len(train_loader), batch_time=batch_time,
loss=losses, top1=top1))
测试函数为
def validate(val_loader, model, criterion, epoch):
"""Perform validation on the validation set"""
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, (input, target) in enumerate(testloader):
target = target.cuda(async=True)
input = input.cuda()
input_var = torch.autograd.Variable(input, volatile=True)
target_var = torch.autograd.Variable(target, volatile=True)
# compute output
output = model(input_var)
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1 = accuracy(output.data, target, topk=(1,))[0]
losses.update(loss.data[0], input.size(0))
top1.update(prec1[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
i, len(val_loader), batch_time=batch_time, loss=losses,
top1=top1))
print(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))
# log to TensorBoard
#if args.tensorboard:
# log_value('val_loss', losses.avg, epoch)
# log_value('val_acc', top1.avg, epoch)
return top1.avg
def accuracy(output, target, topk=(1,)):
"""Computes the precision@k for the specified values of k"""
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0)
res.append(correct_k.mul_(100.0 / batch_size))
return res
最后添加主函数：
model = VGG16()
model = model.cuda()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), 0.001,
momentum=0.9,
nesterov=True,
weight_decay=1*e-4)
for epoch in range(0, 200):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(trainloader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(testloader, model, criterion, epoch)
best_prec1=0
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
print('Best accuracy: ', best_prec1)代码中的文件路径按照自己文件夹做相应修改。