About training LightGlue in two stages

[EATMustard]

Hello, thanks for your work!
In the LightGlue paper, it is mentioned, 'we first train it to predict correspondences and only after train the confidence classifier. The latter thus does not impact the accuracy at the final layer or the convergence of the training.' However, in the actual code, it seems that they are trained together, and there is no distinction between the two stages.

[sarlinpe]

We figured out later that we could simplify the implementation by training the confidence classifier at the same time as the rest of the model but detaching its inputs. This only adds a marginal computational overhead and the end results are strictly identical.

glue-factory/gluefactory/models/matchers/lightglue.py

Lines 590 to 595 in 4a82835

	losses["confidence"] += self.token_confidence[i].loss(
	pred["ref_descriptors0"][:, i],
	pred["ref_descriptors1"][:, i],
	params_i["log_assignment"],
	pred["log_assignment"],
	) / (N - 1)

glue-factory/gluefactory/models/matchers/lightglue.py

Lines 75 to 88 in 4a82835

	def loss(self, desc0, desc1, la_now, la_final):
	logit0 = self.token[0](desc0.detach()).squeeze(-1)
	logit1 = self.token[0](desc1.detach()).squeeze(-1)
	la_now, la_final = la_now.detach(), la_final.detach()
	correct0 = (
	la_final[:, :-1, :].max(-1).indices == la_now[:, :-1, :].max(-1).indices
	)
	correct1 = (
	la_final[:, :, :-1].max(-2).indices == la_now[:, :, :-1].max(-2).indices
	)
	return (
	self.loss_fn(logit0, correct0.float()).mean(-1)
	+ self.loss_fn(logit1, correct1.float()).mean(-1)
	) / 2.0