How can I extend this to multi-label classification?

[JeanHung]

No description provided.

[andsteing]

You would need to extend the script in two places:

First, use something like

label = tf.scatter_nd(labels[:, None], tf.ones(tf.shape(data['label'])[0]), (num_classes,))

to get multiple labels in one-hot format in the input processing

vision_transformer/vit_jax/input_pipeline.py

Line 215 in 85c4f53

label = tf.one_hot(data['label'], num_classes) # pylint: disable=no-value-for-parameter

Second, use a sigmoid loss instead of the cross-entropy loss here:

vision_transformer/vit_jax/train.py

Lines 52 to 58 in 85c4f53

	def loss_fn(params, images, labels):
	logits = apply_fn(
	dict(params=params),
	rngs=dict(dropout=dropout_rng),
	inputs=images,
	train=True)
	return cross_entropy_loss(logits=logits, labels=labels)

[JeanHung]

Thank you so much!