OneFlow nn.Linear权重载入问题

rentainhe commented 3 years ago

本次issue的背景是，在复现mlp-mixer的时候，搭建了和pytorch一样的模型，载入了一样的权重，结果精度为0，通过逐层比较载入后的权重的值发现，在中间的linear层所载入的值和pytorch相反，像是transpose后的值

debug的代码已经提交了相应的draft PR 可以复现：https://github.com/Oneflow-Inc/models/pull/132

相应使用的pytorch权重下载链接：https://oneflow-public.oss-cn-beijing.aliyuncs.com/model_zoo/mlp_mixer/mixer_b16_224.pth

pre-requiresite

pip install pytorch
pip install timm

在compare.py下，比较方法如下：

构建pytorch的mlp-mixer模型，并载入权重后导出state_dict

def torch_weight_load(model, checkpoint_path):
state_dict = torch.load(checkpoint_path)
model.load_state_dict(state_dict)
return model.state_dict()

构建oneflow模型，并转换pytorch的权重载入后导出state_dict

def oneflow_weight_load(model, checkpoint_path):
state_dict = torch.load(checkpoint_path)
new_parameters = dict()
for key, value in state_dict.items():
    if "num_batches_tracked" not in key:
        val = value.detach().cpu().numpy()
        new_parameters[key] = val.astype(np.float32)
model.load_state_dict(new_parameters)
return model.state_dict()

逐层比较weight的值，打印不一致的key

def compare(torch_dict, of_dict):
for torch_key, of_key in zip(torch_dict, of_dict):
    value = torch_dict[torch_key].numpy() - of_dict[of_key].numpy()
    if value.mean() != 0.0:
        print(torch_key)

通过打印的结果发现在中间的linear层都不一致: J5`B2MV7HEH422UF__C9V S

发现了以下问题：

中间层的linear权重都不一致
patch_embed的卷积层权重也不一致？
最后一层的head输出是一致的？？（最后一层head也是一层linear）

可能的原因：

我compare代码写错了（可能性比较小因为载入权重的模型我验证过了，结果和pytorch完全不一致，输入变量都控制了，pytorch的结果为69，oneflow为0，权重载入出问题）
之前vit中也有类似的问题，当时的解决办法是在linear载入的时候加入转置
linear层的权重载入可能和torch是反的，需要注意一下转置问题

MARD1NO commented 3 years ago

如果某些层权重像是你说的转置，那我感觉做矩阵乘的时候维度就应该不匹配会报错了。可以在你那个if判断里，print下torch和oneflow权重的shape等信息。
判断weight差异，可以加点 abs_val, abs_mean, abs_var来看下最大差异到多少，哪个层有最大差异
我感觉以 value.mean == 0 来判断权重是否一致有点点不太严谨，可以试试设置一个比较小的tolerance

rentainhe commented 3 years ago

可以，我这边也比较疑惑，在之前对程鹏他们对BERT的时候和我之前对ViT的时候都存在这个问题，但是当时是因为linear层的shape的h和w一致，可能就是一个转置就可以解决了，这边mixer里的是(192, 384)，我也很疑惑是怎么在h和w不同的shape里还能出现载入的问题

判断差异我这边是之前每个层的权重相减的结果都打印出来了，如果权重是一样的话打印里显示的都是0，我后续再增加一些别的判断吧，因为载入的权重数据类型什么的都是一致的，也是从同一个torch权重里转出来的，我转成numpy应该也是一致的吧

我之前还做过比较是，打印linear层中torch导出的权重的第一列，和oneflow导出权重的第一行，发现是一致的，说明可能是在载入的时候flatten+reshape才有可能出现这种问题，我再去check一下是什么问题，应该通过reshape操作可以让他们载一样的权重

我这边还有比较疑惑的点是：

第一层卷积参数不一致
中间层linear参数是不一致的
最后一层linear的参数是一致的

这些我今天去做一些更严谨的实验看看

L1aoXingyu commented 3 years ago

我觉得与其调用 model.load_state_dict() 还不如直接 for 循环去便利所有的 key 和 value，然后把判断一下 weight.shape 和 value.shape，然后通过 weight.data.copy_(value.data) 的方式进行 load

rentainhe commented 3 years ago

我觉得与其调用 model.load_state_dict() 还不如直接 for 循环去便利所有的 key 和 value，然后把判断一下 weight.shape 和 value.shape，然后通过 weight.data.copy_(value.data) 的方式进行 load

我进行一下尝试

tea321000 commented 3 years ago

已经定位到问题，是由于module的_load_from_state_dict里的copy()实现有bug导致的，fc层的weight在执行.copy()后返回的param与原来的input_param并不相等。更换成原地的flow.tensor操作即可fix。

                try:
                    with flow.no_grad():
                        param[:] = flow.tensor(input_param, dtype = param.dtype, shape = param.shape, device = param.device)
                        # param.copy_(input_param)
                        print(param.numpy() == input_param)

L1aoXingyu commented 3 years ago

这样改完之后就不需要 transpose 进行 load 了吗

tea321000 commented 3 years ago

这样改完之后就不需要 transpose 进行 load 了吗

这个只是目前的过渡方案，最终还是要fix copy_的实现，目前还没定位到是copy_的哪部分实现导致没对齐

rentainhe commented 3 years ago

重新统计了一下转换错误的权重 mlp-mixer中包含一下几个结构:

stem: Conv2d
mlp_tokens.fc1
mlp_tokens.fc2
mlp_channels.fc1
mlp_channels.fc2
head：一层fc

出现错误权重转换的层为：

stem: Conv2d
mlp_tokens.fc1
mlp_channels.fc1
mlp_channels.fc2

转换:

def convert_weight(of_state_dict):
    for key, value in of_state_dict.items():
        # (384, 196) -> 384*196 -> (196, 384) -> (384, 196)
        # 转换mlp_tokens.fc1的权重
        if 'fc1' in key  and 'weight' in key and 'mlp_tokens' in key:
            H, W = value.data.shape
            value.data.copy_(value.data.flatten().reshape(W, H).transpose(0, 1))
        # 转换 mlp_channels.fc1 和 mlp_channels.fc2 的权重
        elif 'fc' in key and 'weight' in key and 'mlp_channels' in key:
            H, W = value.data.shape
            value.data.copy_(value.data.flatten().reshape(W, H).transpose(0, 1))
    return of_state_dict

def compare(torch_dict, of_dict):
    of_dict = convert_weight(of_dict)
    for torch_key, of_key in zip(torch_dict, of_dict):         
        value = torch_dict[torch_key].numpy() - of_dict[of_key].numpy()
        if value.mean() != 0.0:
            print(torch_key)

经过转换后结果只剩下conv_stem权重没有对齐，其他都和pytorch对齐

daquexian commented 3 years ago

已经定位到问题，是由于module的_load_from_state_dict里的copy()实现有bug导致的，fc层的weight在执行.copy()后返回的param与原来的input_param并不相等。更换成原地的flow.tensor操作即可fix。
                try:
                    with flow.no_grad():
                        param[:] = flow.tensor(input_param, dtype = param.dtype, shape = param.shape, device = param.device)
                        # param.copy_(input_param)
                        print(param.numpy() == input_param)

有没有最小可复现的脚本，我修复一下

daquexian commented 3 years ago

在 https://github.com/Oneflow-Inc/oneflow/pull/6117 里修复，待合并

Oneflow-Inc / models

OneFlow nn.Linear权重载入问题 #133

pre-requiresite