pnnx转换模型，pnnx模型输出与torch 原模型输出不一致

[manmantang]

error log | 日志或报错信息 | ログ

pnnx 导出模型，输出与原模型对不上

model | 模型 | モデル

1. original model

import torch
import torch.nn as nn
import torch.nn.functional as F
import pnnx
import torch
from torchvision import transforms

import numpy as np
import copy

class NetVLAD(nn.Module):
    """NetVLAD layer implementation"""

    def __init__(self, num_clusters=64, dim=512, alpha=100.0, normalize_input=True):
        """
        Args:
            num_clusters : int
                The number of clusters
            dim : int
                Dimension of descriptors
            alpha : float
                Parameter of initialization. Larger value is harder assignment.
            normalize_input : bool
                If true, descriptor-wise L2 normalization is applied to input.
        """
        super(NetVLAD, self).__init__()
        self.num_clusters = num_clusters
        self.dim = dim
        self.alpha = alpha
        self.normalize_input = normalize_input
        self.conv = nn.Conv2d(dim, num_clusters, kernel_size=(1, 1), bias=False)
        self.centroids = nn.Parameter(torch.rand(num_clusters, dim), requires_grad=True)

        self.clsts = None
        self.traindescs = None

    def _init_params(self):
        clstsAssign = self.clsts / np.linalg.norm(self.clsts, axis=1, keepdims=True)
        dots = np.dot(clstsAssign, self.traindescs.T)
        dots.sort(0)
        dots = dots[::-1, :] # sort, descending

        self.alpha = (-np.log(0.01) / np.mean(dots[0,:] - dots[1,:])).item()
        self.centroids.data.copy_(torch.from_numpy(self.clsts))
        self.conv.weight.data.copy_(torch.from_numpy(self.alpha*clstsAssign).unsqueeze(2).unsqueeze(3))

    def custom_expand(self,tensor, n):
        original_size = tensor.size()
        tensor_expanded = torch.zeros((n,) + original_size, dtype=torch.float)
        for i in range(n):
            tensor_expanded[i] = tensor.clone()
        return tensor_expanded

    def forward(self, x):
        N, C = x.shape[:2]
        if self.normalize_input:
            x = F.normalize(x, p=2, dim=1)  # across descriptor dim

        # soft-assignment
        soft_assign = self.conv(x).view(N, self.num_clusters, -1)
        soft_assign = F.softmax(soft_assign, dim=1)

        x_flatten = x.view(N, C, -1)
        # calculate residuals to each clusters in one loop
        x_flatten = self.custom_expand(x_flatten, self.num_clusters)
        # x_flatten = x_flatten.expand(self.num_clusters, -1, -1, -1)

        x_flatten = x_flatten.permute(1, 0, 2, 3)
        cur_param = self.centroids.expand(x_flatten.size(-1), -1, -1).permute(1, 2, 0).unsqueeze(0)
        residual = x_flatten - cur_param
        # print(f"residual size : ",residual.size())

        residual *= soft_assign.unsqueeze(2)
        vlad = residual.sum(dim=-1)

        return vlad

model = NetVLAD()
model.eval()

input = torch.randn(1, 512, 30, 40)
output1 = model(input)
print(f"output1 : ",output1)

opt_model = pnnx.export(model, "export_pnnx_models/netvlad.pt", input)
output2 = opt_model(input)
print(f"output2 : ",output2)

how to reproduce | 复现步骤 | 再現方法

1.运行以上源码即可复现问题； 详细问题描述： 如果运行x_flatten = x_flatten.expand(self.num_clusters, -1, -1, -1)，将三维tensor扩展成四维，pnnx无法成功导出 ；如果自己重写一个x_flatten = self.custom_expand(x_flatten, self.num_clusters)替换掉expand,pnnx模型输出错误。

[manmantang]

[nihui]

已确认 pnnx 在折叠常量过程中没有考虑 slice tensor 的情形