pytorch从训练到部署

说明：

最近梳理了从pytorch训练模型，到模型转换，到使用C程序运行模型的一整套过程，结构简单易用理解。由于代码行数比较多就不一一写到文章中了，只讲每个功能主要的代码贴出来了，相信有一些基础的人就已经能够运行了。

感谢参考内容的作者

关于在Ubuntu系统上配置深度学习开发环境

https://blog.csdn.net/XDH19910113/article/details/111470521
训练模型
代码名：

catdog（原名pytorch-train-test-onnx）
核心代码：

def train(args):
# read data
dataloders, dataset_sizes, class_names = ImageDataset(args)

with open(args.class_file, ‘w’) as f:
for name in class_names:
f.writelines(name + ‘\n’)
# use gpu or not
use_gpu = torch.cuda.is_available()
print(“use_gpu:{}”.format(use_gpu))

# get model
model = SimpleNet()

(images, labels) = next(iter(dataloders[‘train’]))
trainsw.add_graph(model=model, input_to_model=images)

if args.resume:
if os.path.isfile(args.resume):
print((“=> loading checkpoint ‘{}'”.format(args.resume)))
model.load_state_dict(torch.load(args.resume))
else:
print((“=> no checkpoint found at ‘{}'”.format(args.resume)))

if use_gpu:
model = torch.nn.DataParallel(model)
model.to(torch.device(‘cuda’))
else:
model.to(torch.device(‘cpu’))

# 用交叉熵损失函数(define loss function)
criterion = nn.CrossEntropyLoss()

# 梯度下降(Observe that all parameters are being optimized)
optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4)

# Decay LR by a factor of 0.98 every 1 epoch
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)

model = train_model(args=args,
model=model,
criterion=criterion,
optimizer=optimizer_ft,
scheduler=exp_lr_scheduler,
num_epochs=args.num_epochs,
dataset_sizes=dataset_sizes,
use_gpu=use_gpu,
dataloders=dataloders)

torch.save(model.state_dict(), os.path.join(args.save_path, ‘scenejudgment.pth’))

代码说明：

这个代码是从网上下的，觉得代码写的很规范，功能很全，就在这个代码的基础上进行了一系列的操作，基础的东西使用别人造好的轮子
下载地址：

https://download.csdn.net/download/dcrmg/11939795
下载之后文件名：

pytorch-train-test-onnx.rar

class SimpleNet(nn.Module): 网络结构
def train(args): 训练
def test(test_model_path, test_img_path, class_file): 测试
def convert_model_to_ONNX(input_img_size, input_pth_model, output_ONNX):
将pth转为onnx
def readImg(path): 数据增强，这个函数读取完图像后直接做了数据增强
def set_parser(): 设置参数，python很常用的一个包

问题说明：
问题1：

这个工程数据有一些问题，太乱了，分类猫狗，长啥样的都有，还有很多错的，因此准确率一直上不去，90%是个坎，后期可以尝试更换一批数据集重新进行尝试
问题2：
报错：

Unexpected key(s) in state_dict: “module.conv1.0.weight”, “module.conv1.0.bias”

1

原因：

训练的时候使用GPU进行训练，会在保存模型时增加“module.”字符串，等到test的时候一张图像一张图像进行测试，不启用GPU，因此读取的模型名前没有“module.”字符串，造成此处报错。将模型的参数全部打印出来就会发现，key值不同，多了module
解决办法：

https://blog.csdn.net/qq_32998593/article/details/89343507

个人使用了其中的办法一
说明：

造成这个错误的代码是test部分的：

model.load_state_dict(state_dict)

1

虽然源代码在这个部分有明确的提示：

if torch.cuda.device_count() > 1:
# 如果有多个GPU，将模型并行化，用DataParallel来操作。这个过程会将key值加一个”module. ***”。
model = nn.DataParallel(model)
model.load_state_dict(checkpoint)