NVIDIA / flownet2-pytorch

Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks
Other
3.08k stars 738 forks source link

Training with MultiScale loss #276

Open zinuok opened 1 year ago

zinuok commented 1 year ago

Hello, I'm trying to fine-tune the whole FlowNet2 network from the provided pre-trained weights, with the multi-scale loss by:

...
--loss=MultiScale \
--loss_norm=L1 \
...

However, the model returns only single flow output from the last fusion layer

        return flownetfusion_flow

and in the MultiScale class, the type of the flow output, which is passed as an argument, is checked

        if type(output) is tuple:
            target = self.div_flow * target
            for i, output_ in enumerate(output):
                target_ = self.multiScales[i](target)
                epevalue += self.loss_weights[i]*EPE(output_, target_)
                lossvalue += self.loss_weights[i]*self.loss(output_, target_)
            return [lossvalue, epevalue]
        else:
            epevalue += EPE(output, target)
            lossvalue += self.loss(output, target)
            return  [lossvalue, epevalue]

The problem is that the type of model output is not always a tuple, so it falls to the else branch, which is the same as a simple L1loss.

Looking at the related issue: issue 28, I guess the multiscale loss was designed to train only a part of the model, i.e. FlowNetC/S, separately. Is it right?

Thank you