目录
原始数据描述
多数原始数据结构如下:即训练数据和测试数据分别放在两个文件夹中(train_filelist、val_filelist); 训练数据和测试数据的标签分别放在两个txt文件中(train.txt、val.txt)
train_filelist文件储存内容为各种类别花朵图片:
val_filelist文件储存内容为各种类别花朵图片:
train.txt 储存的是train_filelist里面各种类别花的分类标签:
如image_06753.jpg 的花的图片类别为 0
如image_07086.jpg 的花的图片类别为 9
val.txt 储存的是val_filelist里面各种类别花的分类标签:
如image_06739.jpg 的花的图片类别为 0
如image_07895.jpg 的花的图片类别为 100
数据集DataLoader制作
如何自定义数据集:
- 先搞定数据和标签的目录结构(得知道到哪里读取数据)
- 写好读取数据和标签路径得函数(根据自己数据集情况来写)
- 完成单个数据与标签读取函数(给dataloader举一个例子)
以花朵数据为例子:
- train.txt以及val.txt 指定了数据路径和标签 (大多数数据集都是这样)
- 这次的任务就是在txt文件中获取图像路径和标签,然后把它们交给dataloader
- 核心代码:按照对应格式传递需要得数据和标签即可
导包
import os
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision
#pip install torchvision
from torchvision import transforms, models, datasets
#https://pytorch.org/docs/stable/torchvision/index.html
import imageio
import time
import warnings
import random
import sys
import copy
import json
from PIL import Image
读取txt文件中的路径和标签
def load_annotations(ann_file): #ann_file 表示要输入的txt文件,如:ann_file = './flower_data/train.txt'
data_infos = {}
with open(ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()] #按行读取txt文件,并列按照空格分开 分开后大概如下:samples = [(image_06734.jpg,0),(image_0875.jpg,99),......]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos
print(load_annotations('./flower_data/train.txt'))
得到的结果部分如下:
分别把数据和标签存到不同list中
因为dataloader到时候会在list中分别读取数据和标签,因此先把数据和标签分别保存在list中
img_label = load_annotations('./flower_data/train.txt')
image_name = list(img_label.keys())
label = list(img_label.values())
获取数据以及标签完整路径
data_dir = './flower_data/'
train_dir = data_dir + '/train_filelist'
valid_dir = data_dir + '/val_filelist'
image_path = [os.path.join(train_dir,img) for img in image_name]
image_path
transform预处理(数据增强)
data_transforms = {
'train':
transforms.Compose([
transforms.Resize(64),
transforms.RandomRotation(45),#随机旋转,-45到45度之间随机选
transforms.CenterCrop(64),#从中心开始裁剪
transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转 选择一个概率概率
transforms.RandomVerticalFlip(p=0.5),#随机垂直翻转
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),#参数1为亮度,参数2为对比度,参数3为饱和度,参数4为色相
transforms.RandomGrayscale(p=0.025),#概率转换成灰度率,3通道就是R=G=B
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#均值,标准差
]),
'valid':
transforms.Compose([
transforms.Resize(64),
transforms.CenterCrop(64),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
设置DataLoader实例化函数
from torch.utils.data import Dataset, DataLoader
class FlowerDataset(Dataset):
def __init__(self, root_dir, ann_file, transform=None):
#root_dir:表示训练集或者测试集数据的路径,如:./flower_data/train_filelist
#ann_file:表示训练集或者测试集标签的路径,如:'./flower_data/train.txt
#transform:表示数据预处理,如数据增强
self.ann_file = ann_file
self.root_dir = root_dir
self.img_label = self.load_annotations()
self.img = [os.path.join(self.root_dir,img) for img in list(self.img_label.keys())] #储存训练集或者测试集的数据,用list装起来
self.label = [label for label in list(self.img_label.values())] #储存训练集或者测试集的标签,用list装起来
self.transform = transform
def __len__(self):
return len(self.img)
def __getitem__(self, idx):
image = Image.open(self.img[idx])
label = self.label[idx]
if self.transform:
image = self.transform(image)
label = torch.from_numpy(np.array(label))
return image, label
def load_annotations(self):
data_infos = {}
with open(self.ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos
实例化DataLoader
train_dataset = FlowerDataset(root_dir=train_dir, ann_file = './flower_data/train.txt', transform=data_transforms['train'])
val_dataset = FlowerDataset(root_dir=valid_dir, ann_file = './flower_data/val.txt', transform=data_transforms['valid'])
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)
试一试数据和标签是否对应
image, label = iter(train_loader).next()
sample = image[0].squeeze()
sample = sample.permute((1, 2, 0)).numpy()
sample *= [0.229, 0.224, 0.225]
sample += [0.485, 0.456, 0.406]
plt.imshow(sample)
plt.show()
print('Label is: {}'.format(label[0].numpy()))
image, label = iter(val_loader).next()
sample = image[0].squeeze()
sample = sample.permute((1, 2, 0)).numpy()
sample *= [0.229, 0.224, 0.225]
sample += [0.485, 0.456, 0.406]
plt.imshow(sample)
plt.show()
print('Label is: {}'.format(label[0].numpy()))
自定义DataLoader完整代码
#导包
import os
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import torch
from torch import nn
import torch.optim as optim
import torchvision
#pip install torchvision
from torchvision import transforms, models, datasets
#https://pytorch.org/docs/stable/torchvision/index.html
import imageio
import time
import warnings
import random
import sys
import copy
import json
from PIL import Image
from torch.utils.data import Dataset, DataLoader
#读取txt文件中的路径和标签
def load_annotations(ann_file): #ann_file 表示要输入的txt文件,如:ann_file = './flower_data/train.txt'
data_infos = {}
with open(ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()] #按行读取txt文件,并列按照空格分开 分开后大概如下:samples = [(image_06734.jpg,0),(image_0875.jpg,99),......]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos
print(load_annotations('./flower_data/train.txt'))
#分别把数据和标签储存到不同的list中
img_label = load_annotations('./flower_data/train.txt')
image_name = list(img_label.keys())
label = list(img_label.values())
#获取数据以及标签的完整路径
data_dir = './flower_data/'
train_dir = data_dir + '/train_filelist'
valid_dir = data_dir + '/val_filelist'
image_path = [os.path.join(train_dir,img) for img in image_name]
#transform预处理(数据增强)
data_transforms = {
'train':
transforms.Compose([
transforms.Resize(64),
transforms.RandomRotation(45),#随机旋转,-45到45度之间随机选
transforms.CenterCrop(64),#从中心开始裁剪
transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转 选择一个概率概率
transforms.RandomVerticalFlip(p=0.5),#随机垂直翻转
transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),#参数1为亮度,参数2为对比度,参数3为饱和度,参数4为色相
transforms.RandomGrayscale(p=0.025),#概率转换成灰度率,3通道就是R=G=B
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])#均值,标准差
]),
'valid':
transforms.Compose([
transforms.Resize(64),
transforms.CenterCrop(64),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
#定义Dataloader实例化函数
class FlowerDataset(Dataset):
def __init__(self, root_dir, ann_file, transform=None):
#root_dir:表示训练集或者测试集数据的路径,如:./flower_data/train_filelist
#ann_file:表示训练集或者测试集标签的路径,如:'./flower_data/train.txt
#transform:表示数据预处理,如数据增强
self.ann_file = ann_file
self.root_dir = root_dir
self.img_label = self.load_annotations()
self.img = [os.path.join(self.root_dir,img) for img in list(self.img_label.keys())] #储存训练集或者测试集的数据,用list装起来
self.label = [label for label in list(self.img_label.values())] #储存训练集或者测试集的标签,用list装起来
self.transform = transform
def __len__(self):
return len(self.img)
def __getitem__(self, idx):
image = Image.open(self.img[idx])
label = self.label[idx]
if self.transform:
image = self.transform(image)
label = torch.from_numpy(np.array(label))
return image, label
def load_annotations(self):
data_infos = {}
with open(self.ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos
#实例化DataLoader
train_dataset = FlowerDataset(root_dir=train_dir, ann_file = './flower_data/train.txt', transform=data_transforms['train'])
val_dataset = FlowerDataset(root_dir=valid_dir, ann_file = './flower_data/val.txt', transform=data_transforms['valid'])
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)
模型建模
#载入实例化好的dataloader数据
dataloaders = {'train':train_loader,'valid':val_loader}
#选模型
model_name = 'resnet' #可选的比较多 ['resnet', 'alexnet', 'vgg', 'squeezenet', 'densenet', 'inception']
#是否用人家训练好的特征来做
feature_extract = True
# 是否用GPU训练
train_on_gpu = torch.cuda.is_available()
if not train_on_gpu:
print('CUDA is not available. Training on CPU ...')
else:
print('CUDA is available! Training on GPU ...')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#选用模型
model_ft = models.resnet18()
model_ft
#模型输出层修改
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, 102))
input_size = 64
model_ft
# 优化器设置
optimizer_ft = optim.Adam(model_ft.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)#学习率每7个epoch衰减成原来的1/10
criterion = nn.CrossEntropyLoss()
#定义训练函数
def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False, filename='best.pth'):
since = time.time()
best_acc = 0
model.to(device)
val_acc_history = []
train_acc_history = []
train_losses = []
valid_losses = []
LRs = [optimizer.param_groups[0]['lr']]
best_model_wts = copy.deepcopy(model.state_dict())
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# 训练和验证
for phase in ['train', 'valid']:
if phase == 'train':
model.train() # 训练
else:
model.eval() # 验证
running_loss = 0.0
running_corrects = 0
# 把数据都取个遍
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
# 清零
optimizer.zero_grad()
# 只有训练的时候计算和更新梯度
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
#print(loss)
# 训练阶段更新权重
if phase == 'train':
loss.backward()
optimizer.step()
# 计算损失
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(dataloaders[phase].dataset)
epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
time_elapsed = time.time() - since
print('Time elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
# 得到最好那次的模型
if phase == 'valid' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
state = {
'state_dict': model.state_dict(),#字典里key就是各层的名字,值就是训练好的权重
'best_acc': best_acc,
'optimizer' : optimizer.state_dict(),#优化器的状态信息
}
torch.save(state, filename)
if phase == 'valid':
val_acc_history.append(epoch_acc)
valid_losses.append(epoch_loss)
scheduler.step(epoch_loss)#学习率衰减
if phase == 'train':
train_acc_history.append(epoch_acc)
train_losses.append(epoch_loss)
print('Optimizer learning rate : {:.7f}'.format(optimizer.param_groups[0]['lr']))
LRs.append(optimizer.param_groups[0]['lr'])
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# 训练完后用最好的一次当做模型最终的结果,等着一会测试
model.load_state_dict(best_model_wts)
return model, val_acc_history, train_acc_history, valid_losses, train_losses, LRs
#模型训练
model_ft, val_acc_history, train_acc_history, valid_losses, train_losses, LRs = train_model(model_ft, dataloaders, criterion, optimizer_ft, num_epochs=20, filename='best.pth')
模型部署(flask)
模型训练好后,需要进一步将模型部署到服务端,以供客户端使用。
模型训练好后,需构造结构如下:
其中flower_data 为花的数据
best.pth 为建模好的模型
flask_server.py 为部署在服务端的代码
flask_predict.py 为部署在客户端的代码
模型服务端部署(flask)
进一步写 flask_server.py 的代码
#导包
import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#from torchvision import transforms as T
from torchvision import transforms, models, datasets
from torch.autograd import Variable
# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False
# 加载模型进来
def load_model():
"""Load the pre-trained model, you can use your model just as easily.
"""
global model
#这里我们直接加载官方工具包里提供的训练好的模型(代码会自动下载)括号内参数为是否下载模型对应的配置信息
model = models.resnet18()
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(nn.Linear(num_ftrs, 102)) # 类别数自己根据自己任务来
#print(model)
checkpoint = torch.load('best.pth')
model.load_state_dict(checkpoint['state_dict'])
#将模型指定为测试格式
model.eval()
#是否使用gpu
if use_gpu:
model.cuda()
#数据预处理(目的是将客户端输入进来的图片预处理成目标的图片格式,也即做模型建模时候的图片格式)
def prepare_image(image, target_size):
"""Do image preprocessing before prediction on any data.
:param image: original image
:param target_size: target image size
:return:
preprocessed image
"""
#针对不同模型,image的格式不同,但需要统一至RGB格式
if image.mode != 'RGB':
image = image.convert("RGB")
# Resize the input image and preprocess it.(按照所使用的模型将输入图片的尺寸修改,并转为tensor)
image = transforms.Resize(target_size)(image)
image = transforms.ToTensor()(image)
# Convert to Torch.Tensor and normalize. mean与std (RGB三通道)这里的参数和数据集中是对应的,训练过程中一致
image = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image) #这个要和模型建模的时候设置的参数一致
# Add batch_size axis.增加一个维度,用于按batch测试 本次这里一次测试一张
image = image[None]
if use_gpu:
image = image.cuda()
return Variable(image, volatile=True) #不需要求导
# 开启服务 这里的predict只是一个名字,可自定义
@app.route("/predict", methods=["POST"])
def predict():
# Initialize the data dictionary that will be returned from the view.
#做一个标志,刚开始无图像传入时为false,传入图像时为true
data = {"success": False}
# 如果收到请求
if flask.request.method == 'POST':
#判断是否为图像
if flask.request.files.get("image"):
# Read the image in PIL format
# 将收到的图像进行读取
image = flask.request.files["image"].read()
image = Image.open(io.BytesIO(image)) #二进制数据
# 利用上面的预处理函数将读入的图像进行预处理
image = prepare_image(image, target_size=(64, 64))
preds = F.softmax(model(image), dim=1)
results = torch.topk(preds.cpu().data, k=3, dim=1)
results = (results[0].cpu().numpy(), results[1].cpu().numpy())
#将data字典增加一个key,value,其中value为list格式
data['predictions'] = list()
# Loop over the results and add them to the list of returned predictions
for prob, label in zip(results[0][0], results[1][0]):
#label_name = idx2label[str(label)]
r = {"label": str(label), "probability": float(prob)}
#将预测结果添加至data字典
data['predictions'].append(r)
# Indicate that the request was a success.
data["success"] = True
# 将最终结果以json格式文件传出
return flask.jsonify(data)
"""
test_json = {
"status_code": 200,
"success": {
"message": "image uploaded",
"code": 200
},
"video":{
"video_name":opt['source'].split('/')[-1],
"video_path":opt['source'],
"description":"1",
"length": str(hour)+','+str(minute)+','+str(round(second,4)),
"model_object_completed":model_flag
}
"status_txt": "OK"
}
response = requests.post(
'http://xxx.xxx.xxx.xxx:8090/api/ObjectToKafka/',,
data={'json': str(test_json)})
"""
if __name__ == '__main__':
print("Loading PyTorch model and Flask starting server ...")
print("Please wait until server has fully started")
#先加载模型
load_model()
#再开启服务
app.run(port='5012')
运行flask_server.py 代码后(开启服务端),结果如下:开辟了红框中的一个端口
模型客户端构造
进一步写 flask_predict.py 的代码,以便给客户端发送请求
import requests
import argparse
# url和端口 这个端口要和开启后的服务端显示的端口一致
flask_url = 'http://127.0.0.1:5012/predict'
def predict_result(image_path):
#啥方法都行
image = open(image_path, 'rb').read()
payload = {'image': image}
#request发给server.
r = requests.post(flask_url, files=payload).json()
# 成功的话在返回.
if r['success']:
# 输出结果.
for (i, result) in enumerate(r['predictions']):
print('{}. {}: {:.4f}'.format(i + 1, result['label'],
result['probability']))
# 失败了就打印.
else:
print('Request failed')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Classification demo')
parser.add_argument('--file', default='./flower_data/train_filelist/image_06998.jpg', type=str, help='test image file') #default='./flower_data/train_filelist/image_06998.jpg' 这个就是自己想要预测的照片,可以是任何自己想预测的照片
args = parser.parse_args()
predict_result(args.file)