Hello everyone,
I’m trying to get a script of a model.
Here is the code(I removed most of the code in init function to avoid verbosity):
(spec.matrix is a square matrix that its shape is not predetermined.)
class Cell(nn.Module):
def __init__(self, spec, in_channels, out_channels):
super(Cell, self).__init__()
self.spec = spec
self.num_vertices = self.spec.matrix.shape[0]
# vertex_channels[i] = number of output channels of vertex i
self.vertex_channels = ComputeVertexChannels(in_channels, out_channels, self.spec.matrix)
def forward(self, x: Any) -> Any:
tensors = [x]
out_concat = []
for t in range(1, self.num_vertices - 1):
fan_in = []
for src in range(1, t):
if self.spec.matrix[src, t]:
fan_in.append(Truncate(tensors[src], torch.tensor(self.vertex_channels[t])))
if self.spec.matrix[0, t]:
fan_in.append(self.input_op[t](x))
# perform operation on node
vertex_input = sum(fan_in)
vertex_output = self.vertex_op[t](vertex_input)
tensors.append(vertex_output)
if self.spec.matrix[t, self.num_vertices - 1]:
out_concat.append(tensors[t])
if not out_concat:
assert self.spec.matrix[0, self.num_vertices - 1]
outputs = self.input_op[self.num_vertices - 1](tensors[0])
else:
if len(out_concat) == 1:
outputs = out_concat[0]
else:
outputs = torch.cat(out_concat, 1)
if self.spec.matrix[0, self.num_vertices - 1]:
outputs += self.input_op[self.num_vertices - 1](tensors[0])
return outputs
def Truncate(inputs, channels):
"""Slice the inputs to channels if necessary."""
input_channels = inputs.size()[1]
if input_channels < channels:
raise ValueError('input channel < output channels for truncate')
elif input_channels == channels:
return inputs # No truncation necessary
else:
# Truncation should only be necessary when channel division leads to
# vertices with +1 channels. The input vertex should always be projected to
# the minimum channel count.
assert input_channels - channels == 1
return inputs[:, :channels, :, :]
It returns the following RuntimeError:
Truncate(Tensor inputs, Tensor channels) -> (Tensor):
Expected a value of type 'Tensor (inferred)' for argument 'inputs' but instead found type 'Any'.
Inferred 'inputs' to be of type 'Tensor' because it was not annotated with an explicit type.:
File "model.py", line 136
for src in range(1, t):
if self.spec.matrix[src, t]:
fan_in.append(Truncate(tensors[src], torch.tensor(self.vertex_channels[t])))
~~~~~~~~ <--- HERE
if self.spec.matrix[0, t]:
As error declares, I should annotate the input in tensors. However, doing that will result in RuntimeError:Expected integer literal for index ....
As I understood from this post I tried adding @ torch.jit.interface before the class Cell(nn.Module). It returned RuntimeError:interface declarations must have a return type annotated. and it underlines all the lines in the init function.
Although I can shun the problem by adding a @torch.jit.ignore before the forward function in the Cell Module, it seems a bit of an unreliable solution.
Thank you in advance for any guidance.