BN结构分析

import torch

def print_named_parameters(model):
    for k, (name, param) in enumerate(model.named_parameters()):
        print('[{}] {:<25}: {}'.format(k+1, name, param.shape))
     

def print_named_buffers(model):
    for k, (name, module) in enumerate(model.named_buffers()):
        print('[{}] {:<25}: {}'.format(k+1, name, module.shape))
        

if __name__ == '__main__':
    num_batches, num_channels, height, width = 32, 16, 7, 7
    x = torch.randn(num_batches, num_channels, height, width)
    batchnorm2d = torch.nn.BatchNorm2d(num_channels)
    y = batchnorm2d(x)
    print(y.shape)
    print_named_parameters(batchnorm2d)
    print_named_buffers(batchnorm2d)