python保存中间变量（学习笔记）

无乎648

于 2024-03-27 22:01:29 发布

阅读量510

点赞数 7

分类专栏：图像处理深度学习文章标签： python 学习笔记

本文链接：https://blog.csdn.net/WUHU648/article/details/137090276

版权

深度学习同时被 2 个专栏收录

13 篇文章 6 订阅

订阅专栏

图像处理

7 篇文章 143 订阅

订阅专栏

python保存中间变量

原因：

最近在部署dust3r算法，虽然在本地部署了，也能测试出一定的结果，但是发现无法跑很多图片，为了能够测试多张图片跑出来的模型，于是就在打算在autodl上部署算法，但是由于官方给定的代码是训练好模型后通过可视化三维模型的形式来给出的效果，所以在服务器上没有办法来可视化三维模型（可能有办法，但是总是有解决不了的报错，于是便放弃）

产生思路

打算把官方中的代码分成两部分，上部分是训练好的模型output变量，将output保存下来，下载到本地上，在本地上加载output变量，进而完成后续的代码操作。

保存中间变量的方式

通过下面方式output变量会以output.pkl的文件形式保存在当前文件夹下

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函数实现Python对象的保存。支持自定义的对象，非常方便。Pandas的DataFrame和Obspy的Stream也都可以保存成pickle的格式。主要是以二进制的形式来保存成一种无逻辑的文件。

解决原来的问题

dust3r官方给的代码如下，其中服务器主要是在scene.show()这行代码中无法运行。

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 = 4
    schedule = 'cosine'
    lr = 0.01
    niter = 100

    model = 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/'

    # 获取当前目录下的所有文件
    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 predictions
    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.imgs
    focals = scene.get_focals()
    poses = scene.get_im_poses()
    pts3d = scene.get_pts3d()
    confidence_masks = scene.get_masks()

    # visualize reconstruction
    scene.show()

    # find 2D-2D matches between the two images
    from dust3r.utils.geometry import find_reciprocal_matches, xy_grid
    pts2d_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 matches
    import numpy as np
    from matplotlib import pyplot as pl
    n_viz = 10
    match_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].T
        pl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)
    pl.show(block=True)

将代码分成两部分，上部分由服务器来跑，下部分由本地来跑。

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 = 32
    schedule = 'cosine'
    lr = 0.01
    niter = 300
    model = load_model(model_path, device)
    # load_images can take a list of images or a directory
    base_dir = 'croco/assets/'
    # 获取当前目录下的所有文件
    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 pickle
    pickle.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 = 1
    schedule = 'cosine'
    lr = 0.01
    niter = 300
    base_dir = 'croco/assets/'
    # 获取当前目录下的所有文件
    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 pickle
    f = 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.imgs
    focals = scene.get_focals()
    poses = scene.get_im_poses()
    pts3d = scene.get_pts3d()
    confidence_masks = scene.get_masks()
    # visualize reconstruction
    scene.show()
    # find 2D-2D matches between the two images
    from dust3r.utils.geometry import find_reciprocal_matches, xy_grid
    pts2d_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 matches
    import numpy as np
    from matplotlib import pyplot as pl
    n_viz = 10
    match_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].T
        pl.plot([x0, x1 + W0], [y0, y1], '-+', color=cmap(i / (n_viz - 1)), scalex=False, scaley=False)
    pl.show(block=True)