Writing logic (loops and conditionals) inside the loss function

Hello there,

I am working on a project where the output of the network is a minibatch of 12 channel tensors of size 500x500. In the loss i want to take a each of these tensors and access each element through it arbirtrarily(I will have the information of which coordinate I will go to). The reason I mention arbitrary is because I cannot use slicing or any such thing to parallelize it (I guess…). Is there any way I can do it? I know if I manually loop through the elements one by one its gonna slow down my training immensely. Any idea what should I look into?

Thanks in advance!


Could you share a code samples that shows exactly what are the inputs (and which one requires grad) and what you expect as output?

hey thanks for responding. Yes I can show in code but I believe I can explain better here first, if thats okay.

So its a dataset of multiview images. I am feeding in 4 images at a time of the same scene from 3 different views. The network gives out 4 tensors for each of those 4 images. Right, in the loss I want to minimize two things. One is the simple MSE with the ground truth tensors. Alongside that I want to add one more loss to it. So I have the mapping where pixel 1 of image 1 lies in image 2, image 3 and image 4. So I want to go make sure the network predicts the same thing over the 3 images. More specifically if pixel(0, 0) in image 1 is (10, 10) in image 2, (20,20) in image 3 and (45, 45) in image 4 then for the 4 out put tensors values should be equal over those 4 indices. Otherwise some loss value is gonna get added alongside the MSE.

Also I would like to mention I am doing regression not classification. Each of those tensors are of size 500x500x13. So for each pixel you have 13 values, each ranging from [0-1]. The input is RGB image sized 500x500x3.


So the element-wise mse should be fairly simple to do.

For the specific pixels, given the coordinates, you can use a function like .gather() to get the values at the given coordinates, then you can apply your prefered loss functions to these values.

If I have the indices stored in a file as list, so it will be like a text file full of lists where each list consists of 4 pair values (x1, y1), (x2, y2), (x3, y3), (x4, y4). So I can read these 4 pairs from the files use each of them as indices for each of thensor using the gather method and then apply the loss?

Sorry for being naive, but will this create any severe speed issues due to the file access?

I would load these things in a similar way you load any other sample for your model.
But if these are saved in a file as plain numbers, I would suggest. If it is small enough, do a preprocessing step that reads the whole file, fill in a single Tensor and save that on disk. Then you can just load that Tensor and use it.
Otherwise, store this info next to your input images and load them at the same time as the input image.

In both cases, this should be part of your Dataset so that the Dataloader will reduce any delay in loading it during training.