#
Video Vision Transformer (ViViT)

##
개요

Vivit 모델은 Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, Cordelia Schmid가 제안한 논문 ViViT: A Video Vision Transformer에서 소개되었습니다. 이 논문은 비디오 이해를 위한 pure-transformer 기반의 모델 집합 중에서 최초로 성공한 모델 중 하나를 소개합니다.

논문의 초록은 다음과 같습니다:

우리는 이미지 분류에서 최근 성공을 거둔 순수 트랜스포머 기반 모델을 바탕으로 비디오 분류를 위한 모델을 제안합니다. 본 모델은 입력 비디오로부터 시공간 토큰을 추출한 후, 이를 일련의 트랜스포머 레이어로 인코딩합니다. 비디오에서 발생하는 긴 토큰 시퀀스를 처리하기 위해, 입력의 공간 및 시간 차원을 분리하는 여러 효율적인 모델 변형을 제안합니다. 트랜스포머 기반 모델은 대규모 학습 데이터셋에서만 효과적이라는 것이 일반적이지만, 우리는 학습 중 모델을 효과적으로 정규화하고, 사전 학습된 이미지 모델을 활용함으로써 상대적으로 작은 데이터셋에서도 학습할 수 있는 방법을 보여줍니다. 또한, 철저한 소거(ablation) 연구를 수행하고 Kinetics 400 및 600, Epic Kitchens, Something-Something v2, Moments in Time을 포함한 여러 비디오 분류 벤치마크에서 최첨단 성과를 달성하여, 기존의 3D 합성곱 신경망 기반 방법들을 능가합니다.

이 모델은 jegormeister가 기여하였습니다. 원본 코드(JAX로 작성됨)는 여기에서 확인할 수 있습니다.

##
VivitConfig

###
class transformers.VivitConfig

< source >

( image_size = 224 num_frames = 32 tubelet_size = [2, 16, 16] num_channels = 3 hidden_size = 768 num_hidden_layers = 12 num_attention_heads = 12 intermediate_size = 3072 hidden_act = 'gelu_fast' hidden_dropout_prob = 0.0 attention_probs_dropout_prob = 0.0 initializer_range = 0.02 layer_norm_eps = 1e-06 qkv_bias = True **kwargs )

Parameters

- image_size (`int`, optional, defaults to 224) — The size (resolution) of each image.
- num_frames (`int`, optional, defaults to 32) — The number of frames in each video.
- tubelet_size (`List[int]`, optional, defaults to `[2, 16, 16]`) — The size (resolution) of each tubelet.
- num_channels (`int`, optional, defaults to 3) — The number of input channels.
- hidden_size (`int`, optional, defaults to 768) — Dimensionality of the encoder layers and the pooler layer.
- num_hidden_layers (`int`, optional, defaults to 12) — Number of hidden layers in the Transformer encoder.
- num_attention_heads (`int`, optional, defaults to 12) — Number of attention heads for each attention layer in the Transformer encoder.
- intermediate_size (`int`, optional, defaults to 3072) — Dimensionality of the “intermediate” (i.e., feed-forward) layer in the Transformer encoder.
- hidden_act (`str` or `function`, optional, defaults to `"gelu_fast"`) — The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`, `"relu"`, `"selu"`, `"gelu_fast"` and `"gelu_new"` are supported.
- hidden_dropout_prob (`float`, optional, defaults to 0.0) — The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
- attention_probs_dropout_prob (`float`, optional, defaults to 0.0) — The dropout ratio for the attention probabilities.
- initializer_range (`float`, optional, defaults to 0.02) — The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
- layer_norm_eps (`float`, optional, defaults to 1e-06) — The epsilon used by the layer normalization layers.
- qkv_bias (`bool`, optional, defaults to `True`) — Whether to add a bias to the queries, keys and values.

This is the configuration class to store the configuration of a VivitModel. It is used to instantiate a ViViT model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the ViViT google/vivit-b-16x2-kinetics400 architecture.

Configuration objects inherit from PretrainedConfig and can be used to control the model outputs. Read the documentation from PretrainedConfig for more information.

Example:

    >>> from transformers import VivitConfig, VivitModel

    >>> # Initializing a ViViT google/vivit-b-16x2-kinetics400 style configuration
    >>> configuration = VivitConfig()

    >>> # Initializing a model (with random weights) from the google/vivit-b-16x2-kinetics400 style configuration
    >>> model = VivitModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config

##
VivitImageProcessor

###
class transformers.VivitImageProcessor

< source >

( do_resize: bool = True size: typing.Dict[str, int] = None resample: Resampling = <Resampling.BILINEAR: 2> do_center_crop: bool = True crop_size: typing.Dict[str, int] = None do_rescale: bool = True rescale_factor: typing.Union[int, float] = 0.00784313725490196 offset: bool = True do_normalize: bool = True image_mean: typing.Union[float, typing.List[float], NoneType] = None image_std: typing.Union[float, typing.List[float], NoneType] = None **kwargs )

Parameters

- do_resize (`bool`, optional, defaults to `True`) — Whether to resize the image’s (height, width) dimensions to the specified `size`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
- size (`Dict[str, int]` optional, defaults to `{"shortest_edge" -- 256}`): Size of the output image after resizing. The shortest edge of the image will be resized to `size["shortest_edge"]` while maintaining the aspect ratio of the original image. Can be overriden by `size` in the `preprocess` method.
- resample (`PILImageResampling`, optional, defaults to `Resampling.BILINEAR`) — Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the `preprocess` method.
- do_center_crop (`bool`, optional, defaults to `True`) — Whether to center crop the image to the specified `crop_size`. Can be overridden by the `do_center_crop` parameter in the `preprocess` method.
- crop_size (`Dict[str, int]`, optional, defaults to `{"height" -- 224, "width": 224}`): Size of the image after applying the center crop. Can be overridden by the `crop_size` parameter in the `preprocess` method.
- do_rescale (`bool`, optional, defaults to `True`) — Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale` parameter in the `preprocess` method.
- rescale_factor (`int` or `float`, optional, defaults to `1/127.5`) — Defines the scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the `preprocess` method.
- offset (`bool`, optional, defaults to `True`) — Whether to scale the image in both negative and positive directions. Can be overriden by the `offset` in the `preprocess` method.
- do_normalize (`bool`, optional, defaults to `True`) — Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess` method.
- image_mean (`float` or `List[float]`, optional, defaults to `IMAGENET_STANDARD_MEAN`) — Mean to use if normalizing the image. This is a float or list of floats the length of the number of channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
- image_std (`float` or `List[float]`, optional, defaults to `IMAGENET_STANDARD_STD`) — Standard deviation to use if normalizing the image. This is a float or list of floats the length of the number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.

Constructs a Vivit image processor.

####
preprocess

< source >

( videos: typing.Union[ForwardRef('PIL.Image.Image'), numpy.ndarray, ForwardRef('torch.Tensor'), list['PIL.Image.Image'], list[numpy.ndarray], list['torch.Tensor']] do_resize: bool = None size: typing.Dict[str, int] = None resample: Resampling = None do_center_crop: bool = None crop_size: typing.Dict[str, int] = None do_rescale: bool = None rescale_factor: float = None offset: bool = None do_normalize: bool = None image_mean: typing.Union[float, typing.List[float], NoneType] = None image_std: typing.Union[float, typing.List[float], NoneType] = None return_tensors: typing.Union[str, transformers.utils.generic.TensorType, NoneType] = None data_format: ChannelDimension = <ChannelDimension.FIRST: 'channels_first'> input_data_format: typing.Union[str, transformers.image_utils.ChannelDimension, NoneType] = None )

Parameters

- videos (`ImageInput`) — Video frames to preprocess. Expects a single or batch of video frames with pixel values ranging from 0 to 255. If passing in frames with pixel values between 0 and 1, set `do_rescale=False`.
- do_resize (`bool`, optional, defaults to `self.do_resize`) — Whether to resize the image.
- size (`Dict[str, int]`, optional, defaults to `self.size`) — Size of the image after applying resize.
- resample (`PILImageResampling`, optional, defaults to `self.resample`) — Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only has an effect if `do_resize` is set to `True`.
- do_center_crop (`bool`, optional, defaults to `self.do_centre_crop`) — Whether to centre crop the image.
- crop_size (`Dict[str, int]`, optional, defaults to `self.crop_size`) — Size of the image after applying the centre crop.
- do_rescale (`bool`, optional, defaults to `self.do_rescale`) — Whether to rescale the image values between `[-1 - 1]` if `offset` is `True`, `[0, 1]` otherwise.
- rescale_factor (`float`, optional, defaults to `self.rescale_factor`) — Rescale factor to rescale the image by if `do_rescale` is set to `True`.
- offset (`bool`, optional, defaults to `self.offset`) — Whether to scale the image in both negative and positive directions.
- do_normalize (`bool`, optional, defaults to `self.do_normalize`) — Whether to normalize the image.
- image_mean (`float` or `List[float]`, optional, defaults to `self.image_mean`) — Image mean.
- image_std (`float` or `List[float]`, optional, defaults to `self.image_std`) — Image standard deviation.
- return_tensors (`str` or `TensorType`, optional) — The type of tensors to return. Can be one of:
  - Unset: Return a list of `np.ndarray`.
  - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
  - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
  - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
  - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
- data_format (`ChannelDimension` or `str`, optional, defaults to `ChannelDimension.FIRST`) — The channel dimension format for the output image. Can be one of:
  - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
  - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
  - Unset: Use the inferred channel dimension format of the input image.
- input_data_format (`ChannelDimension` or `str`, optional) — The channel dimension format for the input image. If unset, the channel dimension format is inferred from the input image. Can be one of:
  - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
  - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
  - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.

Preprocess an image or batch of images.

##
VivitModel

###
class transformers.VivitModel

< source >

( config add_pooling_layer = True )

Parameters

- config (VivitConfig) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the from_pretrained() method to load the model weights.

The bare ViViT Transformer model outputting raw hidden-states without any specific head on top. This model is a PyTorch torch.nn.Module subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

####
forward

< source >

( pixel_values: typing.Optional[torch.FloatTensor] = None head_mask: typing.Optional[torch.FloatTensor] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None interpolate_pos_encoding: bool = False return_dict: typing.Optional[bool] = None ) →
transformers.modeling_outputs.BaseModelOutputWithPooling or `tuple(torch.FloatTensor)`

Parameters

- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`) — Pixel values. Pixel values can be obtained using VivitImageProcessor. See VivitImageProcessor.preprocess() for details.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, optional) — Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
  - 1 indicates the head is not masked,
  - 0 indicates the head is masked.
- output_attentions (`bool`, optional) — Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail.
- output_hidden_states (`bool`, optional) — Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail.
- interpolate_pos_encoding (`bool`, optional, `False`) — Whether to interpolate the pre-trained position encodings.
- return_dict (`bool`, optional) — Whether or not to return a ModelOutput instead of a plain tuple.

Returns

transformers.modeling_outputs.BaseModelOutputWithPooling or `tuple(torch.FloatTensor)`

A transformers.modeling_outputs.BaseModelOutputWithPooling or a tuple of `torch.FloatTensor` (if `return_dict=False` is passed or when `config.return_dict=False`) comprising various elements depending on the configuration (VivitConfig) and inputs.

- last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`) — Sequence of hidden-states at the output of the last layer of the model.

- pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`) — Last layer hidden-state of the first token of the sequence (classification token) after further processing through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns the classification token after processing through a linear layer and a tanh activation function. The linear layer weights are trained from the next sentence prediction (classification) objective during pretraining.

- hidden_states (`tuple(torch.FloatTensor)`, optional, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`) — Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.

  Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.

- attentions (`tuple(torch.FloatTensor)`, optional, returned when `output_attentions=True` is passed or when `config.output_attentions=True`) — Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.

  Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

The VivitModel forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

Examples:

    >>> import av
    >>> import numpy as np

    >>> from transformers import VivitImageProcessor, VivitModel
    >>> from huggingface_hub import hf_hub_download

    >>> np.random.seed(0)

    >>> def read_video_pyav(container, indices):
    ...     '''
    ...     Decode the video with PyAV decoder.
    ...     Args:
    ...         container (`av.container.input.InputContainer`): PyAV container.
    ...         indices (`List[int]`): List of frame indices to decode.
    ...     Returns:
    ...         result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3).
    ...     '''
    ...     frames = []
    ...     container.seek(0)
    ...     start_index = indices[0]
    ...     end_index = indices[-1]
    ...     for i, frame in enumerate(container.decode(video=0)):
    ...         if i > end_index:
    ...             break
    ...         if i >= start_index and i in indices:
    ...             frames.append(frame)
    ...     return np.stack([x.to_ndarray(format="rgb24") for x in frames])

    >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
    ...     '''
    ...     Sample a given number of frame indices from the video.
    ...     Args:
    ...         clip_len (`int`): Total number of frames to sample.
    ...         frame_sample_rate (`int`): Sample every n-th frame.
    ...         seg_len (`int`): Maximum allowed index of sample's last frame.
    ...     Returns:
    ...         indices (`List[int]`): List of sampled frame indices
    ...     '''
    ...     converted_len = int(clip_len * frame_sample_rate)
    ...     end_idx = np.random.randint(converted_len, seg_len)
    ...     start_idx = end_idx - converted_len
    ...     indices = np.linspace(start_idx, end_idx, num=clip_len)
    ...     indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
    ...     return indices

    >>> # video clip consists of 300 frames (10 seconds at 30 FPS)
    >>> file_path = hf_hub_download(
    ...     repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset"
    ... )
    >>> container = av.open(file_path)

    >>> # sample 32 frames
    >>> indices = sample_frame_indices(clip_len=32, frame_sample_rate=1, seg_len=container.streams.video[0].frames)
    >>> video = read_video_pyav(container=container, indices=indices)

    >>> image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
    >>> model = VivitModel.from_pretrained("google/vivit-b-16x2-kinetics400")

    >>> # prepare video for the model
    >>> inputs = image_processor(list(video), return_tensors="pt")

    >>> # forward pass
    >>> outputs = model(**inputs)
    >>> last_hidden_states = outputs.last_hidden_state
    >>> list(last_hidden_states.shape)
    [1, 3137, 768]

##
VivitForVideoClassification

###
class transformers.VivitForVideoClassification

< source >

( config )

Parameters

- config (VivitConfig) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the from_pretrained() method to load the model weights.

ViViT Transformer model with a video classification head on top (a linear layer on top of the final hidden state of the [CLS] token) e.g. for Kinetics-400.

Note that it’s possible to fine-tune ViT on higher resolution images than the ones it has been trained on, by setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained position embeddings to the higher resolution.

This model is a PyTorch torch.nn.Module subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

####
forward

< source >

( pixel_values: typing.Optional[torch.FloatTensor] = None head_mask: typing.Optional[torch.FloatTensor] = None labels: typing.Optional[torch.LongTensor] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None interpolate_pos_encoding: bool = False return_dict: typing.Optional[bool] = None ) →
transformers.modeling_outputs.ImageClassifierOutput or `tuple(torch.FloatTensor)`

Parameters

- pixel_values (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`) — Pixel values. Pixel values can be obtained using VivitImageProcessor. See VivitImageProcessor.preprocess() for details.
- head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, optional) — Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
  - 1 indicates the head is not masked,
  - 0 indicates the head is masked.
- output_attentions (`bool`, optional) — Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail.
- output_hidden_states (`bool`, optional) — Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail.
- interpolate_pos_encoding (`bool`, optional, `False`) — Whether to interpolate the pre-trained position encodings.
- return_dict (`bool`, optional) — Whether or not to return a ModelOutput instead of a plain tuple.
- labels (`torch.LongTensor` of shape `(batch_size,)`, optional) — Labels for computing the image classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).

Returns

transformers.modeling_outputs.ImageClassifierOutput or `tuple(torch.FloatTensor)`

A transformers.modeling_outputs.ImageClassifierOutput or a tuple of `torch.FloatTensor` (if `return_dict=False` is passed or when `config.return_dict=False`) comprising various elements depending on the configuration (VivitConfig) and inputs.

- loss (`torch.FloatTensor` of shape `(1,)`, optional, returned when `labels` is provided) — Classification (or regression if config.num_labels==1) loss.

- logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`) — Classification (or regression if config.num_labels==1) scores (before SoftMax).

- hidden_states (`tuple(torch.FloatTensor)`, optional, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`) — Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called feature maps) of the model at the output of each stage.

- attentions (`tuple(torch.FloatTensor)`, optional, returned when `output_attentions=True` is passed or when `config.output_attentions=True`) — Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, sequence_length)`.

  Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

The VivitForVideoClassification forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

Examples:

    >>> import av
    >>> import numpy as np
    >>> import torch

    >>> from transformers import VivitImageProcessor, VivitForVideoClassification
    >>> from huggingface_hub import hf_hub_download

    >>> np.random.seed(0)

    >>> def read_video_pyav(container, indices):
    ...     '''
    ...     Decode the video with PyAV decoder.
    ...     Args:
    ...         container (`av.container.input.InputContainer`): PyAV container.
    ...         indices (`List[int]`): List of frame indices to decode.
    ...     Returns:
    ...         result (np.ndarray): np array of decoded frames of shape (num_frames, height, width, 3).
    ...     '''
    ...     frames = []
    ...     container.seek(0)
    ...     start_index = indices[0]
    ...     end_index = indices[-1]
    ...     for i, frame in enumerate(container.decode(video=0)):
    ...         if i > end_index:
    ...             break
    ...         if i >= start_index and i in indices:
    ...             frames.append(frame)
    ...     return np.stack([x.to_ndarray(format="rgb24") for x in frames])

    >>> def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
    ...     '''
    ...     Sample a given number of frame indices from the video.
    ...     Args:
    ...         clip_len (`int`): Total number of frames to sample.
    ...         frame_sample_rate (`int`): Sample every n-th frame.
    ...         seg_len (`int`): Maximum allowed index of sample's last frame.
    ...     Returns:
    ...         indices (`List[int]`): List of sampled frame indices
    ...     '''
    ...     converted_len = int(clip_len * frame_sample_rate)
    ...     end_idx = np.random.randint(converted_len, seg_len)
    ...     start_idx = end_idx - converted_len
    ...     indices = np.linspace(start_idx, end_idx, num=clip_len)
    ...     indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
    ...     return indices

    >>> # video clip consists of 300 frames (10 seconds at 30 FPS)
    >>> file_path = hf_hub_download(
    ...     repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset"
    ... )
    >>> container = av.open(file_path)

    >>> # sample 32 frames
    >>> indices = sample_frame_indices(clip_len=32, frame_sample_rate=4, seg_len=container.streams.video[0].frames)
    >>> video = read_video_pyav(container=container, indices=indices)

    >>> image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
    >>> model = VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400")

    >>> inputs = image_processor(list(video), return_tensors="pt")

    >>> with torch.no_grad():
    ...     outputs = model(**inputs)
    ...     logits = outputs.logits

    >>> # model predicts one of the 400 Kinetics-400 classes
    >>> predicted_label = logits.argmax(-1).item()
    >>> print(model.config.id2label[predicted_label])
    LABEL_116

< > Update on GitHub