三、训练模型调用及转换
上一章《slomo部分源码分析》为了提高测试速度降低GPU显存的要求对源码进行部分修改。本章实现训练模型的调用及转换为torch.jit模型。
首先通过torch.load()加载上一章训练好的模型，读取两张测试图，并转化为灰度图，把两个转换过的灰度图通过torch.cat()合并后进行预测，同时通过torch.jit保存模型。以下是这部分的代码：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load('D:/project/python/Super-SloMo/data/SuperSloMo.pth') # 加载模型
model = model.to(device)
outmodel = torch.load('D:/project/python/Super-SloMo/data/SuperSloMoOut.pth') # 加载模型
outmodel = outmodel.to(device)
model.eval() # 把模型转为test模式
w, h = (640 // 32) * 32, (360 // 32) * 32
batch = []
img1 = cv2.imread("D:/project/python/Super-SloMo/data/0001.jpg", 0) # 读取要预测的灰度图片
frame1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
frame1 = Image.fromarray(frame1)
frame1 = frame1.resize((w, h), Image.ANTIALIAS)
frame1 = frame1.convert('L')
frame1 = trans_forward(frame1)
batch.append(frame1)
img2 = cv2.imread("D:/project/python/Super-SloMo/data/0002.jpg", 0) # 读取要预测的灰度图片
frame2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
frame2 = Image.fromarray(frame2)
frame2 = frame2.resize((w, h), Image.ANTIALIAS)
frame2 = frame2.convert('L')
frame2 = trans_forward(frame2)
batch.append(frame2)
frame1 = torch.stack(batch[:-1])
frame2 = torch.stack(batch[1:])
i1 = frame1.to(device)
i2 = frame2.to(device)
ix = torch.cat([i1, i2], dim=1)
print(np.array(i1.tolist()).shape)
print(np.array(i2.tolist()).shape)
print(np.array(ix.tolist()).shape)
traced_net = torch.jit.trace(model, ix)
traced_net.save("model.pt")
接下来要处理图像合并的模型。这时需要进行双向光流：
首先将上面预测的结果拆分成f01和f10两个张量。然后对张量进行网格采样，然后进行光流处理。把两张测试图做相同的处理，然后把图像张量、预测张量、中间流张量、中间帧张量合并后把模型导出或者预测。下面是其代码：
flow_out = model(ix)
f01 = flow_out[:, :2, :, :]
f10 = flow_out[:, 2:, :, :]
t = 0.5
temp = -t * (1 - t)
co_eff = [temp, t * t, (1 - t) * (1 - t), temp]
ft0 = co_eff[0] * f01 + co_eff[1] * f10
ft1 = co_eff[2] * f01 + co_eff[3] * f10
gi0ft0 = back_warp(i1, ft0, w, h)
gi1ft1 = back_warp(i2, ft1, w, h)
iy = torch.cat((i1, i2, f01, f10, ft1, ft0, gi1ft1, gi0ft0), dim=1)
io = outmodel(iy)
traced_net = torch.jit.trace(outmodel, iy)
traced_net.save("outmodel.pt")
print("模型序列化导出成功")
ft0f = io[:, :2, :, :] + ft0
ft1f = io[:, 2:4, :, :] + ft1
vt0 = F.sigmoid(io[:, 4:5, :, :])
vt1 = 1 - vt0
gi0ft0f = back_warp(i1, ft0f, w, h)
gi1ft1f = back_warp(i2, ft1f, w, h)
co_eff = [1 - t, t]
ft_p = (co_eff[0] * vt0 * gi0ft0f + co_eff[1] * vt1 * gi1ft1f) / \
(co_eff[0] * vt0 + co_eff[1] * vt1)def back_warp(img, flow, W, H):
# Extract horizontal and vertical flows.
u = flow[:, 0, :, :]
v = flow[:, 1, :, :]
gridX, gridY = np.meshgrid(np.arange(W), np.arange(H))
gridX = torch.tensor(gridX, requires_grad=False, device=device)
gridY = torch.tensor(gridY, requires_grad=False, device=device)
x = gridX.unsqueeze(0).expand_as(u).float() + u
y = gridY.unsqueeze(0).expand_as(v).float() + v
# range -1 to 1
x = 2 * (x / W - 0.5)
y = 2 * (y / H - 0.5)
# stacking X and Y
grid = torch.stack((x, y), dim=3)
# Sample pixels using bilinear interpolation.
imgOut = torch.nn.functional.grid_sample(img, grid)
return imgOut下一章使用C++调用导出的两个训练模型进行中间帧的预测。
本章源码TorchScrip.py：https://download.csdn.net/download/u011736517/12558294