CODEWITHSUNDEEP
pytorchonnx
Fabrice Auzanneau
Fabrice Auzanneau

Reputation: 95

How to generate an onnx file with linear layers using Pytorch

I want to create a network on the basis of the vgg16 network, but adding linear layers (Gemm) just after the conv2d layers, for normalization purpose. After that, I want to export the network in an ONNX file.

The first part seems to work: I took the Pytorch code for generating the vgg16 and modified it as follows

import torch.nn as nn


class VGG(nn.Module):

    def __init__(self, features, num_classes=8, init_weights=True):
        super(VGG, self).__init__()
        self.features = features
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.Linear(4096, 4096),  # New shift layer
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.Linear(4096, 4096),  # New shift layer
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 8),
            nn.Linear(8, 8),        # New shift layer
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x



def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    n = 224
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            n = int(n / 2)
        elif v == 'B':
            layers += [nn.AdaptiveAvgPool2d(n)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            linear = nn.Linear(n,n,True)
            if batch_norm:
                layers += [conv2d, linear, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, linear, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)


cfg = {'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M','B'],
}

def vgg16(**kwargs):
    """VGG 16-layer model (configuration "D")

    """
    model = VGG(make_layers(cfg['D']), **kwargs)
    return model

But when I insert the weights and export to onnx, I see that my linear layers are not referred to as Gemm but as {Transpose + Matmult + Add}

Top layers in the onnx file

The Transpose part is the weights matrix and the Add part is for the biases (which are all 0).

Am I wrong to think that it's possible to do this, or is there a way to get a real Gemm layer here or another way to do this normalization (which is simply multiply all outputs by a single value)?

Upvotes: 1

Views: 1727

Answers (1)

Paige Wang
Paige Wang

Reputation: 41

The input data of nn.Linear here is a 4-D tensor, then torch will export it to {Transpose, MatMul, Add}. Only input is 2-D, the GEMM op will be exported.

You can have to look at the source code of Pytorch for more information.

Upvotes: 4

Related Questions