先簡單說一下整體流程，利用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
)
|