research-ms-loss/ret_benchmark/utils/freeze_bn.py

15 lines
431 B
Python

# Copyright (c) Malong Technologies Co., Ltd.
# All rights reserved.
#
# Contact: github@malong.com
#
# This source code is licensed under the LICENSE file in the root directory of this source tree.
# Batch Norm Freezer
# Note: adds an additional 2% improvement on CUB (on others benchmarks, it brings no effect)
def set_bn_eval(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.eval()