How can I find the code path for a particular function or class in PyTorch?


First of all, PyTorch is a very good library for deep learning researchers. I would like to contribute to the improvement of PyTorch. I have a few questions as below.

I want to easily find the code path for a particular function or class.
For example, there is torch.nn.CrossEntropyLoss which is not defined in torch/nn but is defined in torch/nn/modules/

  1. Is there a way to easily find the code path?

  2. We can import CrossEntropyLoss class using import torch.nn.CrossEntropy or import torch.nn.modues.loss.CrossEntropy. How does the former work?

  3. In, torch.nn.L1Loss and torch.nn.MSELoss do not seem to be implemented. Where are these classes implemented?


It’s cause they are using pythons Intra-package References and those files are the key to follow mapping

Such as for your example here:

1 Like

Goal here I would assume they wanted people to be able to call all these numerous classes or functions simply from torch.nn and not have one massive file containing all of them

@dgriff Thank you very much. Your answers is the answer to question 2.
Do you have any idea about question 1 and 3?

It does answer 1 and 3. That init file is all the module mappings :grin:

grep -ir can also work well :slight_smile:

my bad I see what you mean on question 3 I missed before read to quickly. L1Loss seemed to have name remap to AbsCriterion. so code under AbsCriterion. You will find that the case for these others as well:

name_remap = {
    'TemporalConvolution': 'Conv1d',
    'SpatialDilatedConvolution': 'DilatedConv2d',
    'SpatialMaxUnpooling': 'MaxUnpool2d',
    'SpatialReflectionPadding': 'ReflectionPad2d',
    'SpatialReplicationPadding': 'ReplicationPad2d',
    'VolumetricReplicationPadding': 'ReplicationPad3d',
    'VolumetricMaxUnpooling': 'MaxUnpool3d',
    'SoftMax': 'Softmax',
    'LogSoftMax': 'LogSoftmax',
    'HardTanh': 'Hardtanh',
    'HardShrink': 'Hardshrink',
    'SoftPlus': 'Softplus',
    'SoftShrink': 'Softshrink',
    'MSECriterion': 'MSELoss',
    'AbsCriterion': 'L1Loss',
    'BCECriterion': '_BCELoss',  # TODO: move the glue code into THNN
    'ClassNLLCriterion': 'NLLLoss',
    'DistKLDivCriterion': 'KLDivLoss',
    'SpatialClassNLLCriterion': 'NLLLoss2d',
    'MultiLabelMarginCriterion': 'MultiLabelMarginLoss',
    'MultiMarginCriterion': 'MultiMarginLoss',
    'SmoothL1Criterion': 'SmoothL1Loss',
    'SoftMarginCriterion': 'SoftMarginLoss',
  1. I use pycharm IDE ( and it provides “follow” feature. It’s very convenient.

EDIT: There are many IDEs provide following feature, like VS Code (Although I haven’t used this one.)