I’m interested in implementing the Aux loss model described in this paper. I have attached an image of the model architecture below.
A small description of the model and how it works:
This model is used to classify X-ray images. PA and L are two different views of an X-ray for one patient. Specifically, PA is an X-ray viewed from the back and L the same X-ray but from a side view.
The paper leverages DenseNet121 pretrained models (4 dense blocks) for this architecture. In the image below, a circle represents a dense block.
Here’s the logic/control flow of the model:

If the L image is present but the PA image is missing for a patient, the data should only flow through the L branch (i.e. the DenseNet121 model for L).
If the PA image is present but the L image is missing for a patient, the data should only flow through the PA branch (i.e. the DenseNet121 model for PA).
If both images are present, the data should flow through both branches up to the 3rd dense block. After that, the outputs of each of the 3rd dense blocks should be concatenated and passed through another model (assume it’s linear).

image340×512 21.3 KB



I understand how I could implement and train this model jointly for a batch size of 1. However, the paper claims that they can train this model jointly for a batch size of 8. How could I implement this efficiently in PyTorch for a batch size of 8?