Problems related to pytorch to onnx

These problems are in spectral normalization spectral_norm Meet in .

The first thing I met was torch.mv Operator and torch.dot Problems not supported by operators .

at present pytorch Spectral normalization has been officially realized ：spectral_norm, Which includes torch.mv、 torch.dot operator , turn onnx There will be errors

terms of settlement ： take torch.mv and torch.dot use torch.matmul Instead of , But you may need to change yourself tensor Dimensions .（ adopt unsqueeze And so on. ）

After I solve the above two operators , Be able to run torch.onnx.export function , But when transforming inference ：

RuntimeError: invalid argument 0: Tensors must have same number of dimensions: got 2 and 1

It seems that there is something wrong with the dimension , But I searched for hours and couldn't find the problem .

The solution later was , In turn onnx Before , remove spectral_norm.

Refer to ：https://github.com/pytorch/pytorch/issues/27723

The official has realized how to remove spectral_norm Function of ：

def remove_spectral_norm(module, name='weight'):
    r"""Removes the spectral normalization reparameterization from a module.

    Args:
        module (Module): containing module
        name (str, optional): name of weight parameter

    Example:
        >>> m = spectral_norm(nn.Linear(40, 10))
        >>> remove_spectral_norm(m)
    """
    for k, hook in module._forward_pre_hooks.items():
        if isinstance(hook, SpectralNorm) and hook.name == name:
            hook.remove(module)
            del module._forward_pre_hooks[k]
            break
    else:
        raise ValueError("spectral_norm of '{}' not found in {}".format(
            name, module))

    for k, hook in module._state_dict_hooks.items():
        if isinstance(hook, SpectralNormStateDictHook) and hook.fn.name == name:
            del module._state_dict_hooks[k]
            break

    for k, hook in module._load_state_dict_pre_hooks.items():
        if isinstance(hook, SpectralNormLoadStateDictPreHook) and hook.fn.name == name:
            del module._load_state_dict_pre_hooks[k]
            break
    return module

The specific steps are ：

1. Build the model according to the training model（ At this time, it still contains spectral_norm）, And load pretrained model, This pretrained model contains spectral_norm Related parameters of ：weight_orig、weight_u as well as weight_v.

2. Then use the following function , The input of this function is built model, What is completed is the structure of recursive model , When I meet spectral_norm when , Will call the above remove_spectral_norm remove spectral_norm.

def remove_all_spectral_norm(item):
    if isinstance(item, nn.Module):
        try:
            remove_spectral_norm(item)
        except Exception:
            pass
        
        for child in item.children():  
            remove_all_spectral_norm(child)

    if isinstance(item, nn.ModuleList):
        for module in item:
            remove_all_spectral_norm(module)

    if isinstance(item, nn.Sequential):
        modules = item.children()
        for module in modules:
            remove_all_spectral_norm(module)

3. Finally run torch.onnx.export

Here's a little explanation , Ordinary convolution operations are followed by spectral_norm after , The training parameters will be convoluted weight Will turn into weight_orig、weight_u、weight_v These three categories , That is, these parameters are saved when the model is saved .

Through the above removal operation , From weight_orig、weight_u、weight_v Recover weight.