《Pyramid Stereo Matching Network》--PSMNet 测试代码实现 +如何解决GPU内存不足问题

• 导师给了建议，转战DL方向做立体匹配了。
• PSWNet 是cvpr2018的文章，立体匹配的经典，下面介绍一下我是怎么跑通文章的测试代码，以及踩过的坑。
  -对于复现论文，我也是初学者，希望给更多刚入门的师弟师妹一些帮助。

一.下载github源码

1.一般论文都是有github连接的

2.如何阅读github

• 以下就是github中的目录
  
• 注意到 dataloader 文件夹是针对数据输入预处理的，其实就是怎么 crop 图像，怎么对图像做 BatchNormalize
• models 文件夹就是模型的位置
• main.py 这个文件就是 数据预处理-模型加载-训练-验证-输出结果的过程
• 注意有个 README.md ,这个介绍了github整体架构，如何使用这些文件夹，非常重要！！！！！！

二. 如何阅读README.md

1.文章运行所需环境

2.训练和评估过程，翻译一下这里，简单看看

• 由于我只是跑一下模型的测试看看效果，只需要加载与训练模型即可。下载模型和数据集需要 VPN，建议大家自己花几十块钱买一个，很有用！！！
• 这里按照他说的不解压预训练模型，直接torch.load（）
• 按照他的要求
  

三.代码解析

• 跑了一个测试的代码 Test_img.py
• 这个代码 我用 pycharm IDE没跑通。建议 spyder

1.首先遇到的就是 import argparse 这个库带来的问题

• argparse 是个字典，需要将数据、模型、测试图像、参数与变量关联，但路径需要改成自己的。

parser = argparse.ArgumentParser(description='PSMNet')
parser.add_argument('--KITTI', default='2015',
                    help='KITTI version')
parser.add_argument('--datapath', default='/media/jiaren/ImageNet/data_scene_flow_2015/testing/',
                    help='select model')
parser.add_argument('--loadmodel', default='./trained/pretrained_model_KITTI2015.tar',
                    help='loading model')
parser.add_argument('--leftimg', default= './VO04_L.png',
                    help='load model')
parser.add_argument('--rightimg', default= './VO04_R.png',
                    help='load model')                                      
parser.add_argument('--model', default='stackhourglass',
                    help='select model')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='maxium disparity')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
args = parser.parse_args()

2.下载KITTI2015数据集

• 看到全是英文别慌，下载这个
  

KITTI2015

• PSMNet中提到了两个模型，一个是"basic",一个是“stackhourglass”，这里我们是在上面介绍的字典中加载的。
• 我用了 “stackhourglass” 模型

if args.cuda:
    torch.cuda.manual_seed(args.seed)

if args.model == 'stackhourglass':
    model = stackhourglass(args.maxdisp)
elif args.model == 'basic':
    model = basic(args.maxdisp)
else:
    print('no model')
    *将模型放在GPU中
model = nn.DataParallel(model, device_ids=[0])
model.cuda()

3.加载预训练模型

if args.loadmodel is not None:
    print('load PSMNet')
    state_dict = torch.load(args.loadmodel)
    model.load_state_dict(state_dict['state_dict'])

4. 测试模块

def test(imgL,imgR):
*模型为预测模式
        model.eval()
*将数据放入GPU中
        if args.cuda:
           imgL = imgL.cuda()
           imgR = imgR.cuda()     
*由于是测试，我们使用了预训练模型，不需要反向传播
        with torch.no_grad():
            disp = model(imgL,imgR)
        disp = torch.squeeze(disp)
        pred_disp = disp.data.cpu().numpy()

        return pred_disp

• mian 函数分为两部分，图像处理，要和训练阶段处理方式相同，以及测试

def main():
*图像预处理部分
        normal_mean_var = {'mean': [0.485, 0.456, 0.406],
                            'std': [0.229, 0.224, 0.225]}
        infer_transform = transforms.Compose([transforms.ToTensor(),
                                              transforms.Normalize(**normal_mean_var)])    

        imgL_o = Image.open(args.leftimg).convert('RGB')
        imgR_o = Image.open(args.rightimg).convert('RGB')

        imgL = infer_transform(imgL_o)
        imgR = infer_transform(imgR_o) 
       

        # pad to width and hight to 16 times
        if imgL.shape[1] % 16 != 0:
            times = imgL.shape[1]//16       
            top_pad = (times+1)*16 -imgL.shape[1]
        else:
            top_pad = 0

        if imgL.shape[2] % 16 != 0:
            times = imgL.shape[2]//16                       
            right_pad = (times+1)*16-imgL.shape[2]
        else:
            right_pad = 0    

        imgL = F.pad(imgL,(0,right_pad, top_pad,0)).unsqueeze(0)
        imgR = F.pad(imgR,(0,right_pad, top_pad,0)).unsqueeze(0)

*测试部分
        start_time = time.time()
        pred_disp = test(imgL,imgR)
        print('time = %.2f' %(time.time() - start_time))

        
        if top_pad !=0 or right_pad != 0:
            img = pred_disp[top_pad:,:-right_pad]
        else:
            img = pred_disp
        
        img = (img*256).astype('uint16')
        img = Image.fromarray(img)
        img.save('Test_disparity.png')

四.KITTI数据集用法

• UBUNTU系统不会截图，这里不能放图片，请谅解，读者朋友会截图的可以留言，hhhhhhh…
• 以KITTI2012举例，colored_0 是RGB图像左图，colored_1 是RGB图像右图， image_0是灰度图左图， image_1是灰度图右图。
• 根据程序把左右视图放到根目录下，并改名进行测试。

parser.add_argument('--leftimg', default= './VO04_L.png',
                    help='load model')
parser.add_argument('--rightimg', default= './VO04_R.png',
                    help='load model') 

五.记录一些错误

1. 这里提示代码有些是没有对齐的，从github下载的代码有些是串了位置。

unindent does not match any outer indentation level

2. 使用变量没有定义
   解决方法

local variable xxx referenced before assignment

3. GPU内存不够，需要调小BatchSize,或者杀死一些进程，我出错的原因就是之前跑过的进程没有关掉。

RuntimeError: CUDA out of memory. Tried to allocate 352.00 MiB (GPU 0; 7.80 GiB total capacity; 6.45

• 首先查看内存占用

nvidia-smi

• 这里就是全部进程，可以看到我有个进程占了4G

Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0       984      G   /usr/lib/xorg/Xorg                            35MiB |
|    0      1567      G   /usr/lib/xorg/Xorg                           114MiB |
|    0      1700      G   /usr/bin/gnome-shell                         192MiB |
|    0    337434      G   /usr/lib/xorg/Xorg                            35MiB |
|    0    744440      G   /usr/lib/firefox/firefox                       2MiB |
|    0    744477      G   /usr/lib/firefox/firefox                       2MiB |
|    0    747520      G   /home/yht/anaconda3/bin/python                 2MiB |
|    0    747595      C   /home/yht/anaconda3/bin/python              4515MiB |
|    0    767429      G   /usr/lib/firefox/firefox                       2MiB |
|    0    776208      G   /usr/lib/firefox/firefox                       2MiB |
|    0    776395      G   /usr/lib/firefox/firefox                       2MiB |
+-----------------------------------------------------------------------------+

-如果训练或测试停止，就关掉他
通过PID关闭进程，比如：

kill -9 747595