Function copy_layer_param

SwissArmyTransformer/examples/mae/transform_param.py:25–36

in-place copy from src to dst src and dst should be the same layer type, e.g., both are LayerNorm or both are Linear. Or at least, both have same named_parameters name and shape.

(src, dst)

model = MAE(args, layernorm_epsilon=1e-6)

def copy_layer_param(src, dst):
    """
    in-place copy from src to dst
    src and dst should be the same layer type, e.g., both are LayerNorm or both are Linear.
        Or at least, both have same named_parameters name and shape.
    """
    src_dic = dict(src.named_parameters())
    dst_dic = dict(dst.named_parameters())
    for k in dst_dic:
        assert dst_dic[k].data.shape == src_dic[k].data.shape
        dst_dic[k].data = src_dic[k].data
        assert (dst_dic[k].data == src_dic[k].data).all()

def copy_from_param(src, dst):
    assert src.data.shape == dst.data.shape