python保存中間變量

原因：

最近在部署dust3r算法，雖然在本地部署了，也能測試出一定的結(jié)果，但是發(fā)現(xiàn)無法跑很多圖片，為了能夠測試多張圖片跑出來的模型，于是就在打算在autodl上部署算法，但是由于官方給定的代碼是訓(xùn)練好模型后通過可視化三維模型的形式來給出的效果，所以在服務(wù)器上沒有辦法來可視化三維模型（可能有辦法，但是總是有解決不了的報錯，于是便放棄）

產(chǎn)生思路

打算把官方中的代碼分成兩部分，上部分是訓(xùn)練好的模型output變量，將output保存下來，下載到本地上，在本地上加載output變量，進(jìn)而完成后續(xù)的代碼操作。

保存中間變量的方式

通過下面方式output變量會以output.pkl的文件形式保存在當(dāng)前文件夾下

import pickle
output=1 #這里就是要保存的中間變量
pickle.dump(output, open('output.pkl', 'wb'))

通過下面的方式來讀取剛才保存的output.pkl文件，這樣就可以順利保存下來了

 f = open("output.pkl",'rb')output=pickle.loads(f.read())f.close()

原理

pickle是Python官方自帶的庫，提供dump函數(shù)實現(xiàn)Python對象的保存。支持自定義的對象，非常方便。Pandas的DataFrame和Obspy的Stream也都可以保存成pickle的格式。主要是以二進(jìn)制的形式來保存成一種無邏輯的文件。

解決原來的問題

dust3r官方給的代碼如下，其中服務(wù)器主要是在scene.show()這行代碼中無法運(yùn)行。

import osfrom dust3r.inference import inference, load_model
from dust3r.utils.image import load_images
from dust3r.image_pairs import make_pairs
from dust3r.cloud_opt import global_aligner, GlobalAlignerModeif __name__ == '__main__':model_path = "checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"device = 'cuda'batch_size = 4schedule = 'cosine'lr = 0.01niter = 100model = load_model(model_path, device)# load_images can take a list of images or a directory# base_dir = 'tankandtemples/tankandtemples/intermediate/M60/images/'base_dir = 'croco/assets/'# 獲取當(dāng)前目錄下的所有文件files = [os.path.join(base_dir, file) for file in os.listdir(base_dir)]images = load_images(files, size=512)pairs = make_pairs(images, scene_graph='complete', prefilter=None, symmetrize=True)output = inference(pairs, model, device, batch_size=batch_size)# at this stage, you have the raw dust3r predictionsview1, pred1 = output['view1'], output['pred1']view2, pred2 = output['view2'], output['pred2']scene = global_aligner(output, device=device, mode=GlobalAlignerMode.PointCloudOptimizer)loss = scene.compute_global_alignment(init="mst", niter=niter, schedule=schedule, lr=lr)# retrieve useful values from scene:imgs = scene.imgsfocals = scene.get_focals()poses = scene.get_im_poses()pts3d = scene.get_pts3d()confidence_masks = scene.get_masks()# visualize reconstructionscene.show()# find 2D-2D matches between the two imagesfrom dust3r.utils.geometry import find_reciprocal_matches, xy_gridpts2d_list, pts3d_list = [], []for i in range(2):conf_i = confidence_masks[i].cpu().numpy()pts2d_list.append(xy_grid(*imgs[i].shape[:2][::-1])[conf_i])  # imgs[i].shape[:2] = (H, W)pts3d_list.append(pts3d[i].detach().cpu().numpy()[conf_i])reciprocal_in_P2, nn2_in_P1, num_matches = find_reciprocal_matches(*pts3d_list)print(f'found {num_matches} matches')matches_im1 = pts2d_list[1][reciprocal_in_P2]matches_im0 = pts2d_list[0][nn2_in_P1][reciprocal_in_P2]# visualize a few matchesimport numpy as npfrom matplotlib import pyplot as pln_viz = 10match_idx_to_viz = np.round(np.linspace(0, num_matches-1, n_viz)).astype(int)viz_matches_im0, viz_matches_im1 = matches_im0[match_idx_to_viz], matches_im1[match_idx_to_viz]H0, W0, H1, W1 = *imgs[0].shape[:2], *imgs[1].shape[:2]img0 = np.pad(imgs[0], ((0, max(H1 - H0, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)img1 = np.pad(imgs[1], ((0, max(H0 - H1, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)img = np.concatenate((img0, img1), axis=1)pl.figure()pl.imshow(img)cmap = pl.get_cmap('jet')for i in range(n_viz):(x0, y0), (x1, y1) = viz_matches_im0[i].T, viz_matches_im1[i].Tpl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)pl.show(block=True)

將代碼分成兩部分，上部分由服務(wù)器來跑，下部分由本地來跑。

import os
from dust3r.inference import inference, load_model
from dust3r.utils.image import load_images
from dust3r.image_pairs import make_pairs
from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
if __name__ == '__main__':model_path = "checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"device = 'cuda'batch_size = 32schedule = 'cosine'lr = 0.01niter = 300model = load_model(model_path, device)# load_images can take a list of images or a directorybase_dir = 'croco/assets/'# 獲取當(dāng)前目錄下的所有文件files = [os.path.join(base_dir, file) for file in os.listdir(base_dir)]files_new = []for i in range(0,files.__len__(),10):files_new.append(files[i])images = load_images(files_new, size=512)pairs = make_pairs(images, scene_graph='complete', prefilter=None, symmetrize=True)output = inference(pairs, model, device, batch_size=batch_size)import picklepickle.dump(output, open('output.pkl', 'wb'))

本地代碼

import os
from dust3r.inference import inference, load_model
from dust3r.utils.image import load_images
from dust3r.image_pairs import make_pairs
from dust3r.cloud_opt import global_aligner, GlobalAlignerMode
if __name__ == '__main__':model_path = "checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"device = 'cuda'batch_size = 1schedule = 'cosine'lr = 0.01niter = 300base_dir = 'croco/assets/'# 獲取當(dāng)前目錄下的所有文件files = [os.path.join(base_dir, file) for file in os.listdir(base_dir)]files_new = []for i in range(0,files.__len__(),4):files_new.append(files[i])print(files_new)import picklef = open("output.pkl",'rb')output=pickle.loads(f.read())f.close()view1, pred1 = output['view1'], output['pred1']view2, pred2 = output['view2'], output['pred2']scene = global_aligner(output, device=device, mode=GlobalAlignerMode.PointCloudOptimizer)loss = scene.compute_global_alignment(init="mst", niter=niter, schedule=schedule, lr=lr)# retrieve useful values from scene:imgs = scene.imgsfocals = scene.get_focals()poses = scene.get_im_poses()pts3d = scene.get_pts3d()confidence_masks = scene.get_masks()# visualize reconstructionscene.show()# find 2D-2D matches between the two imagesfrom dust3r.utils.geometry import find_reciprocal_matches, xy_gridpts2d_list, pts3d_list = [], []for i in range(2):conf_i = confidence_masks[i].cpu().numpy()pts2d_list.append(xy_grid(*imgs[i].shape[:2][::-1])[conf_i])  # imgs[i].shape[:2] = (H, W)pts3d_list.append(pts3d[i].detach().cpu().numpy()[conf_i])reciprocal_in_P2, nn2_in_P1, num_matches = find_reciprocal_matches(*pts3d_list)print(f'found {num_matches} matches')matches_im1 = pts2d_list[1][reciprocal_in_P2]matches_im0 = pts2d_list[0][nn2_in_P1][reciprocal_in_P2]# visualize a few matchesimport numpy as npfrom matplotlib import pyplot as pln_viz = 10match_idx_to_viz = np.round(np.linspace(0, num_matches-1, n_viz)).astype(int)viz_matches_im0, viz_matches_im1 = matches_im0[match_idx_to_viz], matches_im1[match_idx_to_viz]H0, W0, H1, W1 = *imgs[0].shape[:2], *imgs[1].shape[:2]img0 = np.pad(imgs[0], ((0, max(H1 - H0, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)img1 = np.pad(imgs[1], ((0, max(H0 - H1, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)img = np.concatenate((img0, img1), axis=1)pl.figure()pl.imshow(img)cmap = pl.get_cmap('jet')for i in range(n_viz):(x0, y0), (x1, y1) = viz_matches_im0[i].T, viz_matches_im1[i].Tpl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)pl.show(block=True)

總結(jié)

這種解決辦法也不是根本解決辦法，雖然比較麻煩，但是還是能將項目跑起來，也是沒有辦法的辦法，在此做一個筆記記錄。