OOM Problem when run "dp test" (batch_size=1)

[Manyi-Yang]

Summary
After training a model with the descriptor "neuron": [40, 80, 160], I am trying to using "dp test" command to check my model, then I got some OOM issues. But when I using this model to run DP-MD simulation, it works well. Also when I reduce the network to [25, 50, 100], everything is fine.

Deepmd-kit version, installation way, input file, running commands, error log, etc.
2.0.0.b0, conda

Parton Input parameters:

        "descriptor": {
            "type":             "se_a",
            "sel":              [128],
            "rcut_smth":        1.00,
            "rcut":             7.00,
            "neuron":           [40, 80, 160],
             "axis_neuron": 16,
            "resnet_dt":        false,
            "seed": 722586222
        },
        "fitting_net": {
            "neuron": [
                240,
                240,
                240,
                240
            ],
            "resnet_dt": true,
            "seed": 711230366
        }
    },

Error Information when run dp test

tensorflow.python.framework.errors_impl.ResourceExhaustedError: 2 root error(s) found.
(0) Resource exhausted: OOM when allocating tensor with shape[25600,128,120] and type double on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
[[{{node load/gradients/filter_type_0/MatMul_4_grad/MatMul_1}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

     [[load/o_virial/_27]]

Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

(1) Resource exhausted: OOM when allocating tensor with shape[25600,128,120] and type double on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
[[{{node load/gradients/filter_type_0/MatMul_4_grad/MatMul_1}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

[amcadmus]

You may start with -n 1 option of dp test, which means test only one frame of the test data systems.

[Manyi-Yang]

Thanks
Yes, it works when I set -n N (N<100)

[njzjz]

@tuoping I have seen OOM errors asked here for several times. It could be added in to the FAQ, or we may catch tensorflow.python.framework.errors_impl.ResourceExhaustedError and give a guide to screen.

[tuoping]

I think catch and guide is a better idea. I wrote one but I can't find a place to test it.

[njzjz]

In some GPU codes, the memory is not allocated by TensorFlow's API, making it hard to catch the error. @denghuilu Is it possible to switch to TensorFlow's API?