这种读取图片的方式用的是torch自带的 ImageFolder,读取的文件夹必须在一个大的子文件下,按类别归好类。就像我现在要区分三个类别。
#perpare data set
#train data
train_data=torchvision.datasets.ImageFolder('F:/eyeDataSet/trainData',transform=transforms.Compose(
[
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor()
]))
print(len(train_data))
train_loader=DataLoader(train_data,batch_size=20,shuffle=True)
然后就是fine tuning自己的网络,在torch中可以对整个网络修改后,训练全部的参数也可以只训练其中的一部分,我这里就只训练最后一个全连接层。torchvision中提供了很多常用的模型,比如resnet ,Vgg,Alexnet等等
# prepare model
mode1_ft_res18=torchvision.models.resnet18(pretrained=True)
for param in mode1_ft_res18.parameters():
param.requires_grad=False
num_fc=mode1_ft_res18.fc.in_features
mode1_ft_res18.fc=torch.nn.Linear(num_fc,3)
定义自己的优化器,注意这里的参数只传入最后一层的
#loss function and optimizer
criterion=torch.nn.CrossEntropyLoss()
#parameters only train the last fc layer
optimizer=torch.optim.Adam(mode1_ft_res18.fc.parameters(),lr=0.001)
然后就可以开始训练了,定义好各种参数。
#start train
#label not one-hot encoder
EPOCH=1
for epoch in range(EPOCH):
train_loss=0.
train_acc=0.
for step,data in enumerate(train_loader):
batch_x,batch_y=data
batch_x,batch_y=Variable(batch_x),Variable(batch_y)
#batch_y not one hot
#out is the probability of eatch class
# such as one sample[-1.1009 0.1411 0.0320],need to calculate the max index
# out shape is batch_size * class
out=mode1_ft_res18(batch_x)
loss=criterion(out,batch_y)
train_loss+=loss.data[0]
# pred is the expect class
#batch_y is the true label
pred=torch.max(out,1)[1]
train_correct=(pred==batch_y).sum()
train_acc+=train_correct.data[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step%14==0:
print('Epoch: ',epoch,'Step',step,
'Train_loss: ',train_loss/((step+1)*20),'Train acc: ',train_acc/((step+1)*20))
测试部分和训练部分类似这里就不一一说明。
这样就完整了对自己网络的训练测试,完整代码如下:
import torch
import numpy as np
import torchvision
from torchvision import transforms,utils
from torch.utils.data import DataLoader
from torch.autograd import Variable
#perpare data set
#train data
train_data=torchvision.datasets.ImageFolder('F:/eyeDataSet/trainData',transform=transforms.Compose(
[
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor()
]))
print(len(train_data))
train_loader=DataLoader(train_data,batch_size=20,shuffle=True)
#test data
test_data=torchvision.datasets.ImageFolder('F:/eyeDataSet/testData',transform=transforms.Compose(
[
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor()
]))
test_loader=DataLoader(test_data,batch_size=20,shuffle=True)
# prepare model
mode1_ft_res18=torchvision.models.resnet18(pretrained=True)
for param in mode1_ft_res18.parameters():
param.requires_grad=False
num_fc=mode1_ft_res18.fc.in_features
mode1_ft_res18.fc=torch.nn.Linear(num_fc,3)
#loss function and optimizer
criterion=torch.nn.CrossEntropyLoss()
#parameters only train the last fc layer
optimizer=torch.optim.Adam(mode1_ft_res18.fc.parameters(),lr=0.001)
#start train
#label not one-hot encoder
EPOCH=1
for epoch in range(EPOCH):
train_loss=0.
train_acc=0.
for step,data in enumerate(train_loader):
batch_x,batch_y=data
batch_x,batch_y=Variable(batch_x),Variable(batch_y)
#batch_y not one hot
#out is the probability of eatch class
# such as one sample[-1.1009 0.1411 0.0320],need to calculate the max index
# out shape is batch_size * class
out=mode1_ft_res18(batch_x)
loss=criterion(out,batch_y)
train_loss+=loss.data[0]
# pred is the expect class
#batch_y is the true label
pred=torch.max(out,1)[1]
train_correct=(pred==batch_y).sum()
train_acc+=train_correct.data[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step%14==0:
print('Epoch: ',epoch,'Step',step,
'Train_loss: ',train_loss/((step+1)*20),'Train acc: ',train_acc/((step+1)*20))
#print('Epoch: ', epoch, 'Train_loss: ', train_loss / len(train_data), 'Train acc: ', train_acc / len(train_data))
# test model
mode1_ft_res18.eval()
eval_loss=0
eval_acc=0
for step ,data in enumerate(test_loader):
batch_x,batch_y=data
batch_x,batch_y=Variable(batch_x),Variable(batch_y)
out=mode1_ft_res18(batch_x)
loss = criterion(out, batch_y)
eval_loss += loss.data[0]
# pred is the expect class
# batch_y is the true label
pred = torch.max(out, 1)[1]
test_correct = (pred == batch_y).sum()
eval_acc += test_correct.data[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
print( 'Test_loss: ', eval_loss / len(test_data), 'Test acc: ', eval_acc / len(test_data))