Get output size of a model

1

Hi there,

is there any way one can figure out the output dimension of a model without passing a sample to it?

For example, I have two network net1 and net2. I want to use another network net3, which maps the concatenation of net1 and net2 as the feature to some label. By defining the net3, I have to specify the input dimension which equals net1.output_size + net2.output_size.

My idea would be net1.modules[-1].out_features + net2.modules[-1].out_features. Is there a more elegant solution？

Thanks!

1 Like
2

I think your approach is valid, if the last layer is an instance of nn.Linear and I’m nor sure, if there is a more elegant approach.

3

Thanks for replying.

The last layer is indeed a nn.Linear. Concretely, I want to implement the RNN Transducer here.

4

If you’re looking for the total activiation size of a network, or activation/output per layer, you can use this function that I wrote:


total_output_elements = 0
def calc_total_activation_size(model, call_the_network_function):
    global total_output_elements
    total_output_elements = 0

    def hook(module, input, output):
        global total_output_elements
        total_output_elements += output.numel()
        
    handle = torch.nn.modules.module.register_module_forward_hook(hook)
    result = call_the_network_function()
    handle.remove()
    return result, total_output_elements