先简单说一下整体流程，利用pytorch训练模型并转化为onnx格式，然后配置好dlinfer，利用cv22infer在cv22平台量化序列化模型，展开推理
1训练模型
1.1处理数据集
参考图片下载地址:
https://www.kaggle.com/c/dogs-vs-cats-redux-kernels-edition/data
1.1.1首先继承写一个继承自dataset的类
#继承了Dataset的类
class
DogCat(data.Dataset):
def
__init__(
self
, root, transform
=
None
, train
=
True
, test
=
False
):
'''root表示用于训练的图片地址，前70%用于训练，后30%用于测试
'''
self
.test
=
test
self
.train
=
train
self
.transform
=
transform
imgs
=
[os.path.join(root, img)
for
img
in
os.listdir(root)]
#imgs是一个list，list中放(编号，图片地址)
if
self
.test:
#如果是test集
imgs
=
sorted
(imgs, key
=
lambda
x:
int
(x.split(
'.'
)[
-
2
].split(
'/'
)[
-
1
]))
else
:
#如果是train集
imgs
=
sorted
(imgs, key
=
lambda
x:
int
(x.split(
'.'
)[
-
2
]))
imgs_num
=
len
(imgs)
if
self
.test:
self
.imgs
=
imgs
else
:
random.shuffle(imgs)
if
self
.train:
self
.imgs
=
imgs[:
int
(
0.7
*
imgs_num)]
#self.imgs表示前百分之七十的图片
else
:
self
.imgs
=
imgs[
int
(
0.7
*
imgs_num):]
#self.imgs表示除了前百分之七十的图片
# 作为迭代器必须有的方法，可以用[]符号读取数据
def
__getitem__(
self
, index):
img_path
=
self
.imgs[index]
if
self
.test:
label
=
int
(
self
.imgs[index].split(
'.'
)[
-
2
].split(
'/'
)[
-
1
])
else
:
label
=
1
if
'dog'
in
img_path.split(
'/'
)[
-
1
]
else
0
# 狗的label设为1，猫的设为0
data
=
Image.
open
(img_path)
data
=
self
.transform(data)
return
data, label
def
__len__(
self
):
return
len
(
self
.imgs)
|
1.1.2把传入的图片集转换为Tensor并且改变格式
# 对数据集训练集的处理
transform_train = transforms.Compose([
transforms.Resize((
256
,
256
)), # 先调整图片大小至256x256
transforms.RandomCrop((
224
,
224
)), # 再随机裁剪到224x224
transforms.RandomHorizontalFlip(), # 随机的图像水平翻转，通俗讲就是图像的左右对调
transforms.ToTensor(), #图片转换为Tensor
transforms.Normalize((
0.485
,
0.456
,
0.406
), (
0.229
,
0.224
,
0.2225
)) # 归一化，数值是用ImageNet给出的数值
])
# 对数据集验证集的处理
transform_val = transforms.Compose([
transforms.Resize((
224
,
224
)),
transforms.ToTensor(),
transforms.Normalize((
0.485
,
0.456
,
0.406
), (
0.229
,
0.224
,
0.225
)),
])
device = torch.device(
'cuda'
if
torch.cuda.is_available()
else
'cpu'
) # 若能使用cuda，则使用cuda
trainset = DogCat(
'/mnt/ssd0/zhangwentao/train'
, transform=transform_train)#这是我图片的地址
valset = DogCat(
'/mnt/ssd0/zhangwentao/train'
, transform=transform_val)
#继承了dataset dataset 一捆有多大 数据顺序 多线程输入，=
0
表示单线程
trainloader = torch.utils.data.DataLoader(trainset, batch_size=
20
, shuffle=True, num_workers=
0
)
valloader = torch.utils.data.DataLoader(valset, batch_size=
20
, shuffle=False, num_workers=
0
)
|
1.2修改网络模型
class
Net(nn.Module):
'''pytorch的resnet18接口的最后一层fc层的输出维度是1000。这明显不符合猫狗大战数据集，
因为猫狗大战数据集是二分类的，所以最后一层fc输出的维度已经是2才对。
因此我们需要对resnet18进行最后一层的修改。
'''
def
__init__(
self
, model):
super
(Net,
self
).__init__()
# 取掉model的后1层
self
.resnet_layer
=
nn.Sequential(
*
list
(model.children())[:
-
1
])
self
.Linear_layer
=
nn.Linear(
512
,
2
)
# 加上一层参数修改好的全连接层
def
forward(
self
, x):
x
=
self
.resnet_layer(x)
x
=
x.view(x.size(
0
),
-
1
)
x
=
self
.Linear_layer(x)
return
x
'''具体参考完整代码，这里仅展示如何使用网络模型
model = Net(resnet18(pretrained=True))
model = model.to(device)#转到gpu上运行
'''
|
1.3开始训练
''' for epoch in range(1):
train(epoch)训练
val(epoch) 验证，算损失函数
'''
#训练epoch次
def
train(epoch):
scheduler.step()
#按照Pytorch的定义是用来更新优化器的学习率的
model.train()
#训练模式
train_acc
=
0.0
#准确率
for
batch_idx, (img, label)
in
enumerate
(trainloader):
image
=
Variable(img.cuda())
#放到gpu上
label
=
Variable(label.cuda())
#放到ppu上
optimizer.zero_grad()
#意思是把梯度置零，也就是把loss关于weight的导数变成0.
out
=
model(image)
#投喂图片
loss
=
criterion(out, label)
#利用交叉熵损失函数算出out和label的差别
loss.backward()
#反向传播
optimizer.step()
train_acc
=
get_acc(out, label)
#获得准确率
print
(
"Epoch:%d [%d|%d] loss:%f acc:%f"
%
(epoch, batch_idx,
len
(trainloader), loss.mean(), train_acc))
#验证准确率
def
val(epoch):
print
(
"\nValidation Epoch: %d"
%
epoch)
model.
eval
()
#进入推理模式
total
=
0
correct
=
0
with torch.no_grad():
for
batch_idx, (img, label)
in
enumerate
(valloader):
image
=
Variable(img.cuda())
label
=
Variable(label.cuda())
out
=
model(image)
_, predicted
=
torch.
max
(out.data,
1
)
total
+
=
image.size(
0
)
correct
+
=
predicted.data.eq(label.data).cpu().
sum
()
print
(
"Acc: %f "
%
((
1.0
*
correct.numpy())
/
total))
|
1.4完整代码
from
torchvision.models.resnet
import
resnet18
import
os
import
random
from
PIL
import
Image
import
torch.utils.data as data
import
numpy as np
import
torchvision.transforms as transforms
import
torch
import
torch.nn as nn
import
torch.optim as optim
from
torch.autograd
import
Variable
from
torch.optim.lr_scheduler
import
*
#继承了Dataset的类
class
DogCat(data.Dataset):
def
__init__(
self
, root, transform
=
None
, train
=
True
, test
=
False
):
'''root表示用于训练的图片地址，前70%用于训练，后30%用于测试
'''
self
.test
=
test
self
.train
=
train
self
.transform
=
transform
imgs
=
[os.path.join(root, img)
for
img
in
os.listdir(root)]
#imgs是一个list，list中放(编号，图片地址)
if
self
.test:
#如果是test集
imgs
=
sorted
(imgs, key
=
lambda
x:
int
(x.split(
'.'
)[
-
2
].split(
'/'
)[
-
1
]))
else
:
#如果是train集
imgs
=
sorted
(imgs, key
=
lambda
x:
int
(x.split(
'.'
)[
-
2
]))
imgs_num
=
len
(imgs)
if
self
.test:
self
.imgs
=
imgs
else
:
random.shuffle(imgs)
if
self
.train:
self
.imgs
=
imgs[:
int
(
0.7
*
imgs_num)]
#self.imgs表示前百分之七十
else
:
self
.imgs
=
imgs[
int
(
0.7
*
imgs_num):]
#self.imgs表示除了前百分之七十
# 作为迭代器必须有的方法，可以用[]符号读取数据
def
__getitem__(
self
, index):
img_path
=
self
.imgs[index]
if
self
.test:
label
=
int
(
self
.imgs[index].split(
'.'
)[
-
2
].split(
'/'
)[
-
1
])
else
:
label
=
1
if
'dog'
in
img_path.split(
'/'
)[
-
1
]
else
0
# 狗的label设为1，猫的设为0
data
=
Image.
open
(img_path)
data
=
self
.transform(data)
return
data, label
def
__len__(
self
):
return
len
(
self
.imgs)
# 对数据集训练集的处理
transform_train
=
transforms.Compose([
transforms.Resize((
256
,
256
)),
# 先调整图片大小至256x256
transforms.RandomCrop((
224
,
224
)),
# 再随机裁剪到224x224
transforms.RandomHorizontalFlip(),
# 随机的图像水平翻转，通俗讲就是图像的左右对调
transforms.ToTensor(),
#图片转换为Tensor
transforms.Normalize((
0.485
,
0.456
,
0.406
), (
0.229
,
0.224
,
0.2225
))
# 归一化，数值是用ImageNet给出的数值
])
# 对数据集验证集的处理
transform_val
=
transforms.Compose([
transforms.Resize((
224
,
224
)),
transforms.ToTensor(),
transforms.Normalize((
0.485
,
0.456
,
0.406
), (
0.229
,
0.224
,
0.225
)),
])
device
=
torch.device(
'cuda'
if
torch.cuda.is_available()
else
'cpu'
)
# 若能使用cuda，则使用cuda
trainset
=
DogCat(
'/mnt/ssd0/zhangwentao/train'
, transform
=
transform_train)
valset
=
DogCat(
'/mnt/ssd0/zhangwentao/train'
, transform
=
transform_val)
#继承了dataset dataset 一捆有多大 数据顺序 多线程输入，=0表示单线程
trainloader
=
torch.utils.data.DataLoader(trainset, batch_size
=
20
, shuffle
=
True
, num_workers
=
0
)
valloader
=
torch.utils.data.DataLoader(valset, batch_size
=
20
, shuffle
=
False
, num_workers
=
0
)
#acc为准确率，该函数计算准确率
def
get_acc(output, label):
total
=
output.shape[
0
]
_, pred_label
=
output.
max
(
1
)
num_correct
=
(pred_label
=
=
label).
sum
().item()
return
num_correct
/
total
#训练epoch次
def
train(epoch):
scheduler.step()
#按照Pytorch的定义是用来更新优化器的学习率的
model.train()
#训练模式
train_acc
=
0.0
#准确率
for
batch_idx, (img, label)
in
enumerate
(trainloader):
image
=
Variable(img.cuda())
#放到gpu上
label
=
Variable(label.cuda())
#放到ppu上
optimizer.zero_grad()
#意思是把梯度置零，也就是把loss关于weight的导数变成0.
out
=
model(image)
#投喂图片
loss
=
criterion(out, label)
#利用交叉熵损失函数算出out和label的差别
loss.backward()
#反向传播
optimizer.step()
train_acc
=
get_acc(out, label)
#获得准确率
print
(
"Epoch:%d [%d|%d] loss:%f acc:%f"
%
(epoch, batch_idx,
len
(trainloader), loss.mean(), train_acc))
#验证准确率
def
val(epoch):
print
(
"\nValidation Epoch: %d"
%
epoch)
model.
eval
()
#进入推理模式
total
=
0
correct
=
0
with torch.no_grad():
for
batch_idx, (img, label)
in
enumerate
(valloader):
image
=
Variable(img.cuda())
label
=
Variable(label.cuda())
out
=
model(image)
_, predicted
=
torch.
max
(out.data,
1
)
total
+
=
image.size(
0
)
correct
+
=
predicted.data.eq(label.data).cpu().
sum
()
print
(
"Acc: %f "
%
((
1.0
*
correct.numpy())
/
total))
class
Net(nn.Module):
''' 继承了nn.Module便于操作
pytorch的resnet18接口的最后一层fc层的输出维度是1000。这明显不符合猫狗大战数据集，
因为猫狗大战数据集是二分类的，所以最后一层fc输出的维度已经是2才对。
因此我们需要对resnet18进行最后一层的修改。
'''
def
__init__(
self
, model):
super
(Net,
self
).__init__()
# 取掉model的后1层
self
.resnet_layer
=
nn.Sequential(
*
list
(model.children())[:
-
1
])
self
.Linear_layer
=
nn.Linear(
512
,
2
)
# 加上一层参数修改好的全连接层
def
forward(
self
, x):
x
=
self
.resnet_layer(x)
x
=
x.view(x.size(
0
),
-
1
)
x
=
self
.Linear_layer(x)
return
x
if
__name__
=
=
'__main__'
:
model
=
Net(resnet18(pretrained
=
True
))
model
=
model.to(device)
#转到gpu上运行
#model.parameters()为当前的网络模型的参数空间，lr为学习率，
optimizer
=
torch.optim.SGD(model.parameters(), lr
=
0.001
, momentum
=
0.9
, weight_decay
=
5e
-
4
)
# 设置训练细节
'''torch.optim是一个实现各种优化算法的包，优化器
优化器就是需要根据网络反向传播的梯度信息来更新网络的参数，以起到降低loss函数计算值的作用
'''
scheduler
=
StepLR(optimizer, step_size
=
3
)
criterion
=
nn.CrossEntropyLoss()
#交叉熵损失函数
for
epoch
in
range
(
1
):
train(epoch)
val(epoch)
#torch.save(model, 'modelcatdog.pth') # 保存模型
model.to(
"cpu"
)
#转换回cpu
model.
eval
()
#注意，必须进入推理模式。详细看pytorch官方文档
dummy_input
=
Variable(torch.randn(
1
,
3
,
224
,
224
))
#设置入口的shape
#把模型转换为onnx模型，
torch.onnx.export(model,
dummy_input,
"ans.onnx"
, input_names
=
[
"input"
],
opset_version
=
9
)
|