Why torch.nn.Linear is split into Transpose and Gemm layers in torch.onnx.export()?

[imai-lm]

I looked into the output of torch.onnx.export() and found that every layers declared as torch.nn.Linear() was split into two layers; Transpose then Gemm. I think it is redundant, because Gemm operator of ONNX has transB attribute, which transposes the second argument.
Why wouldn't you use the attribute and simply translate it to Gemm only?

[apaszke]

cc: @ezyang @houseroad

[ezyang]

Yes, this is a known issue. To fix this we'll need to introduce a little optimization pass that fuses transposes into Gemm operators.

[imai-lm]

I have another question: Why haven't you use FC op instead of Gemm? Because it is experimental?

[imai-lm]

I found onnx/onnx#18, where FC was used. I know Gemm was not implemented in ONNX at that time, but why did you change it?

[fmassa]

I think the plan is (was?) to remove FC onnx/onnx#47 (comment)

[imai-lm]

@fmassa Oh, I see. Thanks.

[imai-lm]

Now I found there's another op named "expand" between Transpose and Gemm... What's this??

[ezyang]

It's a bug. Please try #3325

[imai-lm]

@ezyang Sorry, but I couldn't find your PRed branch ezyang:pr/expand-opt on your repo...

[imai-lm]

@ezyang self resolved. https://github.com/ezyang/pytorch-1/tree/pr/expand-opt

[imai-lm]

@ezyang I have tried it, but the result was that there was Expand instead of expand... Can you prune it away??

[ezyang]

@imai-lm So, IIUC, your exported ONNX model had Expand in it (i.e., it didn't error on expand)? That shouldn't happen. Did you recompile the C++ bits? If so, I'll try to repro tomorrow.

[imai-lm]

@ezyang yes, my exported ONNX model had Expand ops in it, and there was no error. I tried it under different environments (Ubuntu, and macOS) with reinstallation and got the same result. What I tried was the following:

from torch.autograd import Variable
import torch.onnx
import torchvision


dummy_input = Variable(torch.randn(10, 3, 224, 224))
model = torchvision.models.alexnet(pretrained=True)
torch.onnx.export(model, dummy_input, "alexnet.proto", verbose=True)

The result model was:

graph(%1 : Float(10, 3, 224, 224)
      %2 : Float(64, 3, 11, 11)
      %3 : Float(64)
      %4 : Float(192, 64, 5, 5)
      %5 : Float(192)
      %6 : Float(384, 192, 3, 3)
      %7 : Float(384)
      %8 : Float(256, 384, 3, 3)
      %9 : Float(256)
      %10 : Float(256, 256, 3, 3)
      %11 : Float(256)
      %12 : Float(4096, 9216)
      %13 : Float(4096)
      %14 : Float(4096, 4096)
      %15 : Float(4096)
      %16 : Float(1000, 4096)
      %17 : Float(1000)) {
  %19 : UNKNOWN_TYPE = Conv[kernel_shape=[11, 11], strides=[4, 4], pads=[2, 2, 2, 2], dilations=[1, 1], group=1](%1, %2), uses = [[%20.i0]];
  %20 : Float(10, 64, 55, 55) = Add[broadcast=1, axis=1](%19, %3), uses = [%21.i0];
  %21 : Float(10, 64, 55, 55) = Relu(%20), uses = [%22.i0];
  %22 : Float(10, 64, 27, 27) = MaxPool[dilations=[1, 1], kernel_shape=[3, 3], pads=[0, 0], strides=[2, 2]](%21), uses = [%23.i0];
  %24 : UNKNOWN_TYPE = Conv[kernel_shape=[5, 5], strides=[1, 1], pads=[2, 2, 2, 2], dilations=[1, 1], group=1](%22, %4), uses = [[%25.i0]];
  %25 : Float(10, 192, 27, 27) = Add[broadcast=1, axis=1](%24, %5), uses = [%26.i0];
  %26 : Float(10, 192, 27, 27) = Relu(%25), uses = [%27.i0];
  %27 : Float(10, 192, 13, 13) = MaxPool[dilations=[1, 1], kernel_shape=[3, 3], pads=[0, 0], strides=[2, 2]](%26), uses = [%28.i0];
  %29 : UNKNOWN_TYPE = Conv[kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1], dilations=[1, 1], group=1](%27, %6), uses = [[%30.i0]];
  %30 : Float(10, 384, 13, 13) = Add[broadcast=1, axis=1](%29, %7), uses = [%31.i0];
  %31 : Float(10, 384, 13, 13) = Relu(%30), uses = [%32.i0];
  %33 : UNKNOWN_TYPE = Conv[kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1], dilations=[1, 1], group=1](%31, %8), uses = [[%34.i0]];
  %34 : Float(10, 256, 13, 13) = Add[broadcast=1, axis=1](%33, %9), uses = [%35.i0];
  %35 : Float(10, 256, 13, 13) = Relu(%34), uses = [%36.i0];
  %37 : UNKNOWN_TYPE = Conv[kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1], dilations=[1, 1], group=1](%35, %10), uses = [[%38.i0]];
  %38 : Float(10, 256, 13, 13) = Add[broadcast=1, axis=1](%37, %11), uses = [%39.i0];
  %39 : Float(10, 256, 13, 13) = Relu(%38), uses = [%40.i0];
  %40 : Float(10, 256, 6, 6) = MaxPool[dilations=[1, 1], kernel_shape=[3, 3], pads=[0, 0], strides=[2, 2]](%39), uses = [%41.i0];
  %41 : Float(10, 9216) = Reshape[shape=[10, 9216]](%40), uses = [%42.i0];
  %43 : Float(10, 9216), %44 : UNKNOWN_TYPE = Dropout[is_test=1, ratio=0.5](%41), uses = [[%47.i0], []];
  %45 : Float(9216!, 4096!) = Transpose[perm=[1, 0]](%12), uses = [%47.i1];
  %46 : Float(10!, 4096) = Expand[shape=[10, 4096]](%13), uses = [%47.i2];
  %47 : Float(10, 4096) = Gemm[alpha=1, beta=1](%43, %45, %46), uses = [%48.i0];
  %48 : Float(10, 4096) = Relu(%47), uses = [%49.i0];
  %50 : Float(10, 4096), %51 : UNKNOWN_TYPE = Dropout[is_test=1, ratio=0.5](%48), uses = [[%54.i0], []];
  %52 : Float(4096!, 4096!) = Transpose[perm=[1, 0]](%14), uses = [%54.i1];
  %53 : Float(10!, 4096) = Expand[shape=[10, 4096]](%15), uses = [%54.i2];
  %54 : Float(10, 4096) = Gemm[alpha=1, beta=1](%50, %52, %53), uses = [%55.i0];
  %55 : Float(10, 4096) = Relu(%54), uses = [%58.i0];
  %56 : Float(4096!, 1000!) = Transpose[perm=[1, 0]](%16), uses = [%58.i1];
  %57 : Float(10!, 1000) = Expand[shape=[10, 1000]](%17), uses = [%58.i2];
  %58 : Float(10, 1000) = Gemm[alpha=1, beta=1](%55, %56, %57), uses = [%0.i0];
  return (%58);
}

[ezyang]

Reproduced. Looking into it.