修改RVC模型WN层forward方法，解决CoreML转换intimplicit算子错误

intimplicit算子错误源于PyTorch追踪过程中生成的隐式浮点转整数操作，CoreML转换器未实现该算子的转换逻辑。针对RVC的WN层，可通过以下修改避免该问题：

1. 显式处理张量切片的索引 #

问题常出现在ConditionalNorm层的张量切片操作中，Python整数索引在追踪时可能被转换为隐式整数转换算子。修改ConditionalNorm的forward方法，显式使用整数张量作为索引：

class ConditionalNorm(nn.Module):
    # 保留原有初始化代码
    def forward(self, x, g=None):
        x = x.transpose(1, 2)
        if g is not None:
            g = self.proj(g).transpose(1, 2)
            # 显式创建与设备一致的整数张量作为切片索引
            split_idx = torch.tensor(self.channels, dtype=torch.long, device=g.device)
            mean = g[:, :, :split_idx]
            if not self.mean_only:
                std = g[:, :, split_idx:]
        else:
            mean = self.bias
            if not self.mean_only:
                std = self.weight
        x = (x - mean)
        if not self.mean_only:
            x = x * torch.exp(std)
        x = x.transpose(1, 2)
        return x

2. 确保WN层forward中类型一致性 #

在WN层的forward方法中，显式保持输出张量与输入类型一致，避免隐式类型转换：

class WN(nn.Module):
    # 保留原有初始化和remove_weight_norm代码
    def forward(self, x, x_mask, g=None, **kwargs):
        output = x
        for layer in self.layers:
            output = layer(output, x_mask, g)
            # 强制输出与输入张量类型一致
            output = output.type_as(x)
        return output * x_mask

3. 改用TorchScript脚本化而非追踪 #

相较于torch.jit.trace，torch.jit.script能更好地处理Python控制流，减少不必要的隐式算子生成。修改测试代码中的模型导出部分：

# 替换trace为script
scripted_model = torch.jit.script(model)

x = ct.TensorType(name='x', shape=test_x.shape)
x_mask = ct.TensorType(name='x_mask', shape=test_x_mask.shape)
g = ct.TensorType(name='g', shape=test_g.shape)

mlmodel = ct.converters.convert(scripted_model, inputs=[x, x_mask, g])

额外注意事项 #

• 确保x_mask为0/1浮点张量（而非随机值），避免非预期的数值计算触发隐式转换
• 升级CoreMLTools至最新版本（>=7.0），新版本对PyTorch算子的支持更完善

内容的提问来源于stack exchange，提问作者David Geldreich