前言
在进行模型训练时，需要注意，在config文件中的一个参数设置：
config.CUDNN.BENCHMARK = True #表明每次训练网络的时候，cudnn都会寻找最优的优化策略，无法复现之前的训练结果，可以进行反复训练，找到一个最优的结果
config.CUDNN.BENCHMARK = False #可以复现之前的训练结果
code
"""
进行模型转换，torch->onnx->tvm
环境要求:tvm转torch模型必须在pytorch1.4的环境下进行，1.4下可以进行torch->onnx->tvm,
同时也可以进行torch->libtorch转换，但是c++中源码编译的libtorch库要和torch->libtorch转换时使用的pytorch版本一致
由于c++工程中使用的libtorch库是1.0,因此torch->libtorch的过程也要在pytorch1.0环境下进行
"""
import numpy as np
from pynvml import * #pip install nvidia-ml-py3,显存监控
import torch.onnx
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
from PIL import Image
#导入模型
# from models_class.resnet_chejian import * #导入需要转换的模型的网络结构
# from msra_resnet import * #centernet转换
# from models_class.simpleBaseLine.pose_resnet_simpleBaseLine import *
from models_class.pose_resnet import *
from models_class.slowfastnet import *
# 加载torch模型进行测试
def load_torch_model_test(img_path,model_path):
# 初始化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cudnn.benchmark = True
model = get_pose_net()
model = model.cuda()
model.to(device)
# 加载模型权重
model.load_state_dict(torch.load(model_path)['state_dict'])
# 进行torch模型测试
test_trans = transforms.Compose([transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
img = Image.open(img_path).convert('RGB')
img = test_trans(img)
# img = Image.fromarray(np.expand_dims(img, 0)) # 扩张0维度
imgblob = img.unsqueeze(0) # tensor禁止进行维度删除
imgblob = torch.autograd.Variable(imgblob).cuda() # 对输入进行数据格式转换，np.array->tensor.cuda.float
torch.no_grad() # 不进行梯度计算
model.eval() # 使用验证模式
heatmap = model(imgblob)
print(heatmap)
return model
# 将torch模型转为libtorch模型,关键点,单卡训练
def gen_torch_to_libtorch_keypoints(model_path, save_path):
# 初始化模型
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# cudnn.benchmark = True
#加载模型结构
# model = get_pose_net() #加载网络架构，不同的网络此处不同，此处用于关键点网络结构加载。
# 加载的原因是：训练模型的时候使用model.module.state_dict()，仅保存模型参数，因此需要加载模型结构
# model = resnet101(class_num=2) #加载网络架构:slowfastnet
# model.to(device)
# model.cuda()
# # 加载模型权重
#第一种方式：
# model.load_state_dict(torch.load(model_path)['state_dict']) #训练时使用以下形式进行保存模型权重以及相关参数
# 使用原因：
# 保存模型的状态，可以设置一些参数，后续训练可以使用，例如中断后可以继续训练
# state = {'epoch': epoch + 1, # 保存的当前轮数
# 'state_dict': mymodel.state_dict(), # 训练好的参数
# 'optimizer': optimizer.state_dict(), # 优化器参数,为了后续的resume
# 'best_pred': best_pred # 当前最好的精度
# , ...., ...}
# 保存模型到checkpoint.pth.tar
# torch.save(state, ‘checkpoint.pth.tar’)
#加载模型参数
# checkpoint = torch.load('model_best.pth.tar')
# model.load_state_dict(checkpoint['state_dict']) # 模型参数
# optimizer.load_state_dict(checkpoint['optimizer']) # 优化参数
# epoch = checkpoint['epoch'] # epoch，可以用于更新学习率等
#第二种方式：
# model.load_state_dict(torch.load(model_path))
#原因是：
#在训练的时候是已下面的形式进行保存了模型参数：
# 保存模型到checkpoint.pth.tar
# torch.save(model.module.state_dict(), ‘checkpoint.pth.tar’)
# 对应的加载模型方法为(这种方法需要先反序列化模型获取参数字典，因此必须先load模型，再load_state_dict)：
# mymodel.load_state_dict(torch.load(‘checkpoint.pth.tar’))
# 第三种方式：
#model = torch.load(‘checkpoint.pth.tar’)
#原因是：
#在训练的时候是已下面的形式进行保存，将模型参数以及网络结构保存在一起：
# 保存
#torch.save(model,‘checkpoint.pth.tar’)
# 加载
# model = torch.load(‘checkpoint.pth.tar’)
#在进行torch转libtorch时，此种方式可能再设定GPU模式
#model.to(device)
#model.cuda()
# 将torch转为libtorch并保存
#不同的网络输入数据大小维度不同，要注意
# example = torch.rand(1, 3, 56, 56).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[n,c,h,w]
# example = torch.rand(1, 3, 1, 112, 112).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[b,c,n,h,w]
#
# traced_script_module = torch.jit.trace(model, example).eval()
# traced_script_module.save(save_path)
#****************车辆关键点估计pytorch->libtorch转换**************#
# 初始化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cudnn.benchmark = True
#加载车辆关键点估计模型＝simplebaseline
model = get_pose_net('',False) #关键点网络结构加载。
#加载车辆关键点估计权重－simplebaseline
model.load_state_dict(torch.load(model_path)) # 训练时使用以下形式进行保存模型权重以及相关参数
# model.load_state_dict(torch.load(model_path)['state_dict']) # 训练时使用以下形式进行保存模型权重以及相关参数
#以下操作是避免pytorch转libtorch的模型显存增加
for p in model.parameters():
p.requires_grad = False
model.to(device)
model.cuda()
model.eval()
# 将torch转为libtorch并保存
example = torch.rand(1, 3, 256, 192).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[n,c,h,w]
# traced_script_module = torch.jit.trace(model, example).eval()
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(save_path)
# 将torch模型转为libtorch模型,关键点,多卡训练
def gen_torch_to_libtorch_keypoints_multiple_gpu(model_path, save_path):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cudnn.benchmark = True
#加载车辆关键点估计模型＝simplebaseline
model = get_pose_net('',False) #关键点网络结构加载。
#加载权重,并且将权重的名字进行转换 module.conv1.weight->conv1.weight
checkpoint=torch.load(model_path)
new_state_dict = OrderedDict()
for k, v in checkpoint.items():
name = k[7:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict) # 训练时使用以下形式进行保存模型权重以及相关参数
# model.load_state_dict(new_state_dict['state_dict']) # 训练时使用以下形式进行保存模型权重以及相关参数
#以下操作是避免pytorch转libtorch的模型显存增加
for p in model.parameters():
p.requires_grad = False
model.to(device)
model.cuda()
model.eval()
# 将torch转为libtorch并保存
example = torch.rand(1, 3, 256, 192).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[n,c,h,w]
# traced_script_module = torch.jit.trace(model, example).eval()
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(save_path)
#将centernet模型 pth->pt
def gen_torch_to_libtorch_centernet(model_path, save_path):
# 加载torch模型，分类模型
# centernet
num_layers = 10
heads = {'hm': 7, 'reg': 2, 'wh': 2}
head_conv = 64
resnet_spec = {10: (BasicBlock, [1, 1, 1, 1]),
18: (BasicBlock, [2, 2, 2, 2]),
34: (BasicBlock, [3, 4, 6, 3]),
50: (Bottleneck, [3, 4, 6, 3]),
101: (Bottleneck, [3, 4, 23, 3]),
152: (Bottleneck, [3, 8, 36, 3])}
name = 'ResNet_10'
block_class, layers = resnet_spec[num_layers]
model = PoseResNet(block_class, layers, heads, head_conv=head_conv)
state_dict = torch.load(model_path)
model.load_state_dict(state_dict['state_dict'])
#以下操作是避免pytorch转libtorch的模型显存增加
for p in model.parameters():
p.requires_grad = False
model.cuda()
model.eval()
example = torch.rand(1, 3, 608, 608).cuda() # 输入图像大小,自行设定,如果使用tensorrt7.0.0多batch使用时，此处必须设定batch的大小
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(save_path)
print("pt convert ok")
# 将torch模型转为libtorch模型,人体姿态转换
def gen_torch_to_libtorch_keypoints_human(model_path, save_path):
# 初始化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cudnn.benchmark = True
#加载车辆关键点估计模型＝simplebaseline
model = get_pose_net('',False) #关键点网络结构加载。
#加载车辆关键点估计权重－simplebaseline
model.load_state_dict(torch.load(model_path)) # 训练时使用以下形式进行保存模型权重以及相关参数
# model.load_state_dict(torch.load(model_path)['state_dict']) # 训练时使用以下形式进行保存模型权重以及相关参数
#以下操作是避免pytorch转libtorch的模型显存增加
for p in model.parameters():
p.requires_grad = False
model.to(device)
model.cuda()
model.eval()
# 将torch转为libtorch并保存
example = torch.rand(1, 3, 256, 192).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[n,c,h,w]
# traced_script_module = torch.jit.trace(model, example).eval()
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(save_path)
# 将torch模型转为libtorch模型,人体姿态转换
def gen_torch_to_libtorch_class(model_path, save_path):
# 初始化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cudnn.benchmark = True
model = resnet29_cls() #pytorch的网络
state_dict = torch.load(model_path)
model.load_state_dict(state_dict,strict=True)
#以下操作是避免pytorch转libtorch的模型显存增加
for p in model.parameters():
p.requires_grad = False
model.eval().cuda()#必须要在此处进行eval()
# 将torch转为libtorch并保存
example = torch.rand(1, 3, 224, 224).cuda() # 任意设定一个输入，输入数据的大小为网络输入的大小[n,c,h,w]
# traced_script_module = torch.jit.trace(model, example).eval()#此处进行eval()转换的模型可能有误
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save(save_path)
if __name__ == "__main__":
#显存监控
GPU_USE=0
nvmlInit() #初始化
handle = nvmlDeviceGetHandleByIndex(GPU_USE) #获得指定GPU的handle
info_begin = nvmlDeviceGetMemoryInfo(handle) #获得显存信息
#路径设置
img_str ="/data_1/Working/project/torch-onnx-tvm/data/000d81205274490bb1729e4e26226f01-3-1@00.jpg"
model_path = "/data1/Working/project/key_Point_Detect/plate_keypoints/exam/torch1.0/exam-0/coco/pose_resnet_50/256x192_d256x3_adam_lr1e-3/model_best_acc_0.46193601710706533.pth"
libtorch_save_path = model_path.replace('.pth','.pt')
# model_path = "/data1/Working/project/Change_lanes_without_lighting/car_light_class/data/exam-1/pth-101/2020-08-12-15-23-22/clip_len_10frame_sample_rate_1_checkpoint_120.pth.tar"
# libtorch_save_path = "/data1/Working/project/Change_lanes_without_lighting/car_light_class/data/exam-1/pth-101/pt/video_classify-120.pt"
#进行转换
# gen_torch_to_libtorch_keypoints_human(model_path,libtorch_save_path)
# gen_torch_to_libtorch_class(model_path,libtorch_save_path)
gen_torch_to_libtorch_keypoints(model_path,libtorch_save_path)
# #输出显存使用mb
info_end = nvmlDeviceGetMemoryInfo(handle)
print("-"*15+"TORCH GPU MEMORY INFO"+"-"*15)
print(" Memory Total: "+str(info_end.total//(1024**2)))
print(" Memory Free: "+str(info_end.free//(1024**2)))
print(" Memory Used: "+str(info_end.used//(1024**2)-info_begin.used//(1024**2)))
print("-" * 40)
3715
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
