Issue 1:
🚀 Feature

automatically inferring the input size after the first forward pass

Motivation

https://mxnet.apache.org/api/python/docs/tutorials/getting-started/to-mxnet/pytorch.html

In PyTorch you have to specify the input size as the first argument of the Linear object. Apache MXNet provides an extra flexibility to network structure by automatically inferring the input size after the first forward pass.

Pitch

I have to specify input of the hidden layer input. I want to make sequence without specifing input size.

Exp: #for now;

pt_net = pt_nn.Sequential(

pt_nn.Linear(28*28, 256),

pt_nn.ReLU(),

pt_nn.Linear(256, 10))

My request; #input size is know and second layer input size also known

pt_net = pt_nn.Sequential(

pt_nn.Linear(256),

pt_nn.ReLU(),

pt_nn.Linear(10))

Alternatives

Additional context

Issue 2:
Motivation

There’s a lot of cases in framework design where we want to be able to treat components like lego blocks when building various model architectures. The following case studies in PyText are the 90% common cases, not the exceptions, ie. most people build models from these architecture families.

Document classifiers in PyText follow a general pattern of Embedding → Representation → Decoder.

Embedding

Word embedding

BiTransformer

Character embedding

Representation

Encoder

DocNN

LSTM (possibly bidirectional)

Pooling layer

Max pool

Attention

etc

Decoder

Almost always an MLP, needs to have an output dimension of #classes

Within the 4 categories of Embedding/Encoder/Pooling/Decoder though, each boundary depends on the dimensionality of the previous component — the encoder depends on the embedding dimension, the pooling depends on the encoder dimension (including eg. whether the LSTM is bidirectional), and the decoder depends on the pooling layer output dimension (and #classes).

However, from a user standpoint, it is extremely desirable to not have to make sure that these dimensions match up exactly. For instance, if a product team wants to experiment with DocNN vs BiLSTM with self attention and two different sizes of word embeddings, they would need to know how to compute the correct input dimensions for their encoders, pooling layers, and MLPs for each of the 4 cases. If they have dense features, those might also factor into the computation for the MLP layer.

PyText solves this problem by having modules of each of these categories compute and track the appropriate dimensionality internally to their constructor, and then constructing the model and wiring the various blocks together before training. See DocNN representation_dim computation for an example of this, and Model.from_config for where we use this to construct the model.

This solution works for us, but only in the case where we take the responsibility of constructing the model, along with all of its subcomponents, away from the user. If we wanted to eg. allow the user to pass in their own representation layer, we wouldn’t be able to ensure the dimensionality match, the user is therefore responsible for also constructing every other model component. As a direct result of this design, PyText is abysmally hard to use in a notebook, and we built a Config system for managing hyperparameters that essentially boils down to allowing users to partially construct modules, then let us take over and finish building the whole model. See one of our BiLSTM Config classes for an example.

This issue is not specific to PyText. Consider this snippet from the demo example from the AllenNLP website:

EMBEDDING_DIM = 6
HIDDEN_DIM = 6

token_embedding = Embedding(num_embeddings=vocab.get_vocab_size('tokens'),
embedding_dim=EMBEDDING_DIM)
word_embeddings = BasicTextFieldEmbedder({"tokens": token_embedding})
lstm = PytorchSeq2SeqWrapper(torch.nn.LSTM(EMBEDDING_DIM, HIDDEN_DIM, batch_first=True))
model = LstmTagger(word_embeddings, lstm, vocab)

(source - https://allennlp.org/tutorials)

Note that similar to PyText, we need to match up the EMBEDDING_DIM between token_embedding and the LSTM module; we therefore couldn’t construct an LstmTagger object with a default argument for just one of word_embeddings or lstm, because there would be no way to create them with the appropriate dimensions.

Proposal - Lazy Modules

We propose adding lazily inferred dimensions to some core Torch components, such as Linear, ConvNd, and RNN. This syntax should allow constructing these components (and any derived or more complex user components containing them) to instantiate the models without fully defining their dimensions, and have them inferred from the tensors passed in the first time their forward function is called. We propose using -1 to represent these dimensions similar to how -1 is used in reshape now (None has also been suggested, covered in alternatives below). For example

>>> l = nn.Linear(-1, 3)
>>> l(torch.rand(2, 4)).size()
[2, 3]
>>> l(torch.rand(2, 5))
TypeError(...)

This would allow building more complex components that could also infer their input dimensions; consider a simple MLP implementation, much like what PyText uses for its decoders:

class MLP(nn.Module):
def __init__(self, dims: List[int], dropout: float = 0.):
super().__init__()
layers = []
for in_dim, dim in zip([-1] + dims, dims):
layers.extend([
nn.Linear(in_dim, dim),
ReLU(),
nn.LayerNorm(dim),
nn.Dropout(dropout),
])
self.layers = nn.Sequential(*layers)

def forward(self, input: torch.Tensor) -> torch.Tensor:
return self.layers(input)

Now this MLP can be used in any PyText model, regardless of what the output shape of representation/pooling is:

mlp1 = MLP([4, 5, 6])
mlp2 = MLP([4, 5, 6])

mlp1(torch.rand(2, 4)) # finalizes mlp1 as expecting input dimension 4
mlp2(torch.rand(2, 5)) # finalizes mlp2 as expecting input dimension 5

And in the cases of PyText and AllenNLP these could allow interfaces that are much simpler and more modular:

class Model(nn.Module):
def __init__(
self,
embedding = WordEmbedding(),
representation = BiLSTM(),
decoder = MLP([4, 5, 6]),
):
self.embedding = embedding
self.representation = representation
self.decoder = decoder

def forward(self, tokens):
embedded = self.embedding(tokens)
representation = self.representation(embedded)
return self.decoder(representation)

model1 = Model()
model2 = Model(representation=DocNN())
model3 = Model(
embedding=WordEmbedding(embedding_dim=200, vocab_size=10000),
decoder=MLP([2, 50, 10]),
)

Sample implementation

This will be supplied shortly following this proposal as a pull request in torch.nn.lazy, but for now here is the proposed implementation.

from torch import nn

# LazyModuleMeta re-implements the type construction semantics for objects to allow
# a slight variant on syntax. Essentially anything with this metaclass can optionally
# execute a single yield statement during its constructor (normal constructors also work fine).
# If it does yield during construction, then a __lazy_init__ function is populated;
# any code occurring before yield in the constuctor will be called as normal during object creation,
# and any code after yield will instead be deferred to the first call of __lazy_init__.

class LazyInitMeta(type):
def __call__(cls, *args, **kwargs):
if hasattr(cls, '__new__'):
obj = cls.__new__(cls, *args, **kwargs)
else:
obj = object.__new__(cls)

def initialize(obj):
res = obj.__init__(*args, **kwargs)
if isinstance(res, types.GeneratorType):
next(res, None)
def lazy_init(call_args):
try:
res.send(call_args)
except StopIteration:
pass
finally:
obj.__lazy_init__ = None
obj.__lazy_init__ = lazy_init
else:
obj.__lazy_init__ = None

if isinstance(obj, cls):
initialize(obj)
return obj

# Here's a Lazy nn.Module implementation using LazyInitMeta, calling __lazy_init__ before the first
# forward pass.

class LazyModule(nn.Module, metaclass=LazyInitMeta):
def __init__(self):
nn.Module.__init__(self)

def __call__(self, *args, **kwargs):
if self.__lazy_init__:
self.__lazy_init__(call_args=(args, kwargs))
return nn.Module.__call__(self, *args, **kwargs)

# Optionally lazy Linear module, based on the current torch implementation of nn.Linear.

class Linear(nn.Linear, LazyModule):
def __init__(self, in_features, out_features, bias=True):
LazyModule.__init__(self)
self.in_features = in_features
self.out_features = out_features
if bias:
self.bias = nn.Parameter(torch.Tensor(out_features))
else:
self.register_parameter('bias', None)

if in_features == -1:
self.register_parameter('weight', None)
([input], _) = yield # lazy init remainder
in_features = self.in_features = input.size()[-1]

self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
self.reset_parameters()

Drawbacks

Lazy modules don’t work perfectly with everything in torch. In particular, if they’re not yet finalized (ie. -1 was passed to a dimension to be inferred, but forward hasn’t been called yet), the following behaviors might be unintuitive:

torch.save/torch.jit.script won’t work properly

both confirmed to work fine with the sample implementation once initialization finishes

module.apply()

any parameter initialization functions

module.parameters() won’t return all of the model parameters, so

creating an optimizer/scheduler normally, ie torch.optim.Adam(module.parameters()) will potentially have unintended side-effects, as not all of the module’s eventual parameters will optimized

calling module.to() will move the parameters that exist, but once forward is called and new parameters are created, they will be on the default device

this can potentially be solved with some tighter init behavior and unit tests, as there’s a finite number of components that need to be made lazy, and they can usually look at other parameters on the module and assume that users want them to be on the same device

alternatively it may make sense to have a concept of module-level device residency

Alternatives

None instead of -1

Use the above proposal, but pass None to indicate an inferred dimension rather than -1. -1 was chosen for the proposal because it mirrors the semantics of reshape and a few other operators, but None may be more intuitive.

Support -1 (or None) for Tensor dimensions

Instead of implementing this at the Module level with lazy modules, implement this at the Tensor level. Have Tensors be able to be lazy- or partially-initialized, and then be able to be finalized with their final dimensions. This would solve most of the above drawbacks above, as for instance we’d always be making and setting all parameters, and it would also likely simplify its integration as all known useful cases of lazy dimensions eventually boil down to passing them to Tensors. However, it’s less clear what this implementation would look like, for instance how the Tensors would generically be able to infer the parameters.

Lazy Parameters

If this were implemented at the Parameter level, it would require pretty radically increasing the complexity of how Parameters work, but could solve most of the drawbacks of this proposal while also leaving Tensor untouched. We can flesh out this design more if there is significant feedback in that direction.

Support via metaclass vs non-metaclass

The current proposal implementation uses a metaclass for Lazy modules to introduce a novel initialization pattern for lazy components (this same pattern could eg. also be used to solve an analogous problems for Optimizer/Scheduler in the future). However, using metaclasses always comes at a cost in Python, especially for libraries that are being built on-top-of, because multiple inheritance in the case of multiple metaclasses is a problem with no generic solution. In other words, if any libraries built on top of pytorch want to use metaclasses in their own Module subclasses, they’ll need to do non-trivial work to make sure those metaclasses are compatible with any metaclass used in pytorch.

The proposal could likely be modified to, for example, use a decorator on the init function instead of a metaclass.