garmentiq.classification.model_definition

Classification model architectures.

Three architectures are provided: CNN3 and CNN4, two convolutional networks of increasing depth, and tinyViT, a DeiT-tiny vision transformer. Any torch.nn.Module can be used instead; these are simply the ones GarmentIQ ships trained weights for.

  1"""Classification model architectures.
  2
  3Three architectures are provided: `CNN3` and `CNN4`, two convolutional networks of
  4increasing depth, and `tinyViT`, a DeiT-tiny vision transformer. Any `torch.nn.Module`
  5can be used instead; these are simply the ones GarmentIQ ships trained weights for.
  6"""
  7import torch.nn as nn
  8import timm
  9
 10
 11class CNN3(nn.Module):
 12    """
 13    CNN3 is a convolutional neural network designed for image classification with moderate depth.
 14
 15    It consists of three convolutional blocks followed by fully connected layers.
 16
 17    Args:
 18        num_classes (int): Number of output classes for classification.
 19
 20    Attributes:
 21        features (nn.Sequential): Convolutional feature extractor.
 22        classifier (nn.Sequential): Fully connected classifier.
 23    """
 24
 25    def __init__(self, num_classes):
 26        """
 27        Initializes the CNN3 model architecture.
 28
 29        Args:
 30            num_classes (int): Number of target classes.
 31        """
 32        super(CNN3, self).__init__()
 33        self.features = nn.Sequential(
 34            # Block 1
 35            nn.Conv2d(3, 64, kernel_size=3, padding=1),
 36            nn.BatchNorm2d(64),
 37            nn.ReLU(inplace=True),
 38            nn.Conv2d(64, 64, kernel_size=3, padding=1),
 39            nn.BatchNorm2d(64),
 40            nn.ReLU(inplace=True),
 41            nn.MaxPool2d(2, 2),
 42            nn.Dropout(0.25),
 43            # Block 2
 44            nn.Conv2d(64, 128, kernel_size=3, padding=1),
 45            nn.BatchNorm2d(128),
 46            nn.ReLU(inplace=True),
 47            nn.Conv2d(128, 128, kernel_size=3, padding=1),
 48            nn.BatchNorm2d(128),
 49            nn.ReLU(inplace=True),
 50            nn.MaxPool2d(2, 2),
 51            nn.Dropout(0.25),
 52            # Block 3
 53            nn.Conv2d(128, 256, kernel_size=3, padding=1),
 54            nn.BatchNorm2d(256),
 55            nn.ReLU(inplace=True),
 56            nn.Conv2d(256, 256, kernel_size=3, padding=1),
 57            nn.BatchNorm2d(256),
 58            nn.ReLU(inplace=True),
 59            nn.AdaptiveAvgPool2d((4, 6)),
 60        )
 61
 62        self.classifier = nn.Sequential(
 63            nn.Linear(256 * 4 * 6, 512),
 64            nn.BatchNorm1d(512),
 65            nn.ReLU(inplace=True),
 66            nn.Dropout(0.5),
 67            nn.Linear(512, 256),
 68            nn.BatchNorm1d(256),
 69            nn.ReLU(inplace=True),
 70            nn.Dropout(0.3),
 71            nn.Linear(256, num_classes),
 72        )
 73
 74    def forward(self, x):
 75        """
 76        Defines the forward pass of the CNN3 model.
 77
 78        Args:
 79            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
 80
 81        Returns:
 82            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
 83        """
 84        x = self.features(x)
 85        x = x.view(x.size(0), -1)
 86        x = self.classifier(x)
 87        return x
 88
 89
 90class CNN4(nn.Module):
 91    """
 92    CNN4 is a deeper convolutional neural network with four convolutional blocks.
 93
 94    This architecture is designed for improved feature extraction by adding an extra conv block
 95    and adaptive pooling before classification layers.
 96
 97    Args:
 98        num_classes (int): Number of output classes for classification.
 99
100    Attributes:
101        features (nn.Sequential): Convolutional feature extractor with four blocks.
102        classifier (nn.Sequential): Fully connected classifier.
103    """
104
105    def __init__(self, num_classes):
106        """
107        Initializes the CNN4 model architecture.
108
109        Args:
110            num_classes (int): Number of target classes.
111        """
112        super(CNN4, self).__init__()
113        self.features = nn.Sequential(
114            # Block 1
115            nn.Conv2d(3, 64, kernel_size=3, padding=1),
116            nn.BatchNorm2d(64),
117            nn.ReLU(inplace=True),
118            nn.Conv2d(64, 64, kernel_size=3, padding=1),
119            nn.BatchNorm2d(64),
120            nn.ReLU(inplace=True),
121            nn.MaxPool2d(2, 2),
122            nn.Dropout(0.25),
123            # Block 2
124            nn.Conv2d(64, 128, kernel_size=3, padding=1),
125            nn.BatchNorm2d(128),
126            nn.ReLU(inplace=True),
127            nn.Conv2d(128, 128, kernel_size=3, padding=1),
128            nn.BatchNorm2d(128),
129            nn.ReLU(inplace=True),
130            nn.MaxPool2d(2, 2),
131            nn.Dropout(0.25),
132            # Block 3
133            nn.Conv2d(128, 256, kernel_size=3, padding=1),
134            nn.BatchNorm2d(256),
135            nn.ReLU(inplace=True),
136            nn.Conv2d(256, 256, kernel_size=3, padding=1),
137            nn.BatchNorm2d(256),
138            nn.ReLU(inplace=True),
139            nn.MaxPool2d(2, 2),
140            nn.Dropout(0.25),
141            # Block 4 (Additional block for more depth)
142            nn.Conv2d(256, 512, kernel_size=3, padding=1),
143            nn.BatchNorm2d(512),
144            nn.ReLU(inplace=True),
145            nn.Conv2d(512, 512, kernel_size=3, padding=1),
146            nn.BatchNorm2d(512),
147            nn.ReLU(inplace=True),
148            # Replace MaxPool with adaptive pooling to fix feature map size
149            nn.AdaptiveAvgPool2d((4, 6)),
150        )
151
152        self.classifier = nn.Sequential(
153            nn.Linear(512 * 4 * 6, 1024),
154            nn.BatchNorm1d(1024),
155            nn.ReLU(inplace=True),
156            nn.Dropout(0.5),
157            nn.Linear(1024, 512),
158            nn.BatchNorm1d(512),
159            nn.ReLU(inplace=True),
160            nn.Dropout(0.3),
161            nn.Linear(512, num_classes),
162        )
163
164    def forward(self, x):
165        """
166        Defines the forward pass of the CNN4 model.
167
168        Args:
169            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
170
171        Returns:
172            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
173        """
174        x = self.features(x)
175        x = x.view(x.size(0), -1)
176        x = self.classifier(x)
177        return x
178
179
180class tinyViT(nn.Module):
181    """
182    tinyViT is a vision transformer model based on the DeiT tiny architecture, pretrained on ImageNet.
183
184    It replaces the original classification head with a linear layer matching the number of classes.
185
186    Args:
187        num_classes (int): Number of output classes.
188        img_size (int): Input image size (assumed square).
189        patch_size (int): Size of the image patches used by the transformer.
190
191    Attributes:
192        backbone (nn.Module): The underlying vision transformer model with modified head.
193    """
194
195    def __init__(
196        self, 
197        num_classes: int, 
198        img_size: int, 
199        patch_size: int, 
200        pretrained: bool = False,
201    ):
202        """
203        Initializes the tinyViT model with a pretrained DeiT backbone.
204
205        Args:
206            num_classes (int): Number of target classes.
207            img_size (int): Input image size (height and width).
208            patch_size (int): Patch size for the transformer.
209            pretrained (bool): Whether to use the pretrained model, deault to `False`.
210        """
211        super(tinyViT, self).__init__()
212        self.backbone = timm.create_model(
213            "deit_tiny_patch16_224",
214            pretrained=pretrained,
215            img_size=img_size,
216            patch_size=patch_size,
217        )
218        in_features = self.backbone.head.in_features
219        self.backbone.head = nn.Linear(in_features, num_classes)
220
221    def forward(self, x):
222        """
223        Defines the forward pass of the tinyViT model.
224
225        Args:
226            x (torch.Tensor): Input tensor of shape (batch_size, 3, img_size, img_size).
227
228        Returns:
229            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
230        """
231        return self.backbone(x)
class CNN3(torch.nn.modules.module.Module): 
12class CNN3(nn.Module):
13    """
14    CNN3 is a convolutional neural network designed for image classification with moderate depth.
15
16    It consists of three convolutional blocks followed by fully connected layers.
17
18    Args:
19        num_classes (int): Number of output classes for classification.
20
21    Attributes:
22        features (nn.Sequential): Convolutional feature extractor.
23        classifier (nn.Sequential): Fully connected classifier.
24    """
25
26    def __init__(self, num_classes):
27        """
28        Initializes the CNN3 model architecture.
29
30        Args:
31            num_classes (int): Number of target classes.
32        """
33        super(CNN3, self).__init__()
34        self.features = nn.Sequential(
35            # Block 1
36            nn.Conv2d(3, 64, kernel_size=3, padding=1),
37            nn.BatchNorm2d(64),
38            nn.ReLU(inplace=True),
39            nn.Conv2d(64, 64, kernel_size=3, padding=1),
40            nn.BatchNorm2d(64),
41            nn.ReLU(inplace=True),
42            nn.MaxPool2d(2, 2),
43            nn.Dropout(0.25),
44            # Block 2
45            nn.Conv2d(64, 128, kernel_size=3, padding=1),
46            nn.BatchNorm2d(128),
47            nn.ReLU(inplace=True),
48            nn.Conv2d(128, 128, kernel_size=3, padding=1),
49            nn.BatchNorm2d(128),
50            nn.ReLU(inplace=True),
51            nn.MaxPool2d(2, 2),
52            nn.Dropout(0.25),
53            # Block 3
54            nn.Conv2d(128, 256, kernel_size=3, padding=1),
55            nn.BatchNorm2d(256),
56            nn.ReLU(inplace=True),
57            nn.Conv2d(256, 256, kernel_size=3, padding=1),
58            nn.BatchNorm2d(256),
59            nn.ReLU(inplace=True),
60            nn.AdaptiveAvgPool2d((4, 6)),
61        )
62
63        self.classifier = nn.Sequential(
64            nn.Linear(256 * 4 * 6, 512),
65            nn.BatchNorm1d(512),
66            nn.ReLU(inplace=True),
67            nn.Dropout(0.5),
68            nn.Linear(512, 256),
69            nn.BatchNorm1d(256),
70            nn.ReLU(inplace=True),
71            nn.Dropout(0.3),
72            nn.Linear(256, num_classes),
73        )
74
75    def forward(self, x):
76        """
77        Defines the forward pass of the CNN3 model.
78
79        Args:
80            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
81
82        Returns:
83            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
84        """
85        x = self.features(x)
86        x = x.view(x.size(0), -1)
87        x = self.classifier(x)
88        return x

CNN3 is a convolutional neural network designed for image classification with moderate depth.

It consists of three convolutional blocks followed by fully connected layers.

Arguments:
  • num_classes (int): Number of output classes for classification.
Attributes:
  • features (nn.Sequential): Convolutional feature extractor.
  • classifier (nn.Sequential): Fully connected classifier.
CNN3(num_classes) 
26    def __init__(self, num_classes):
27        """
28        Initializes the CNN3 model architecture.
29
30        Args:
31            num_classes (int): Number of target classes.
32        """
33        super(CNN3, self).__init__()
34        self.features = nn.Sequential(
35            # Block 1
36            nn.Conv2d(3, 64, kernel_size=3, padding=1),
37            nn.BatchNorm2d(64),
38            nn.ReLU(inplace=True),
39            nn.Conv2d(64, 64, kernel_size=3, padding=1),
40            nn.BatchNorm2d(64),
41            nn.ReLU(inplace=True),
42            nn.MaxPool2d(2, 2),
43            nn.Dropout(0.25),
44            # Block 2
45            nn.Conv2d(64, 128, kernel_size=3, padding=1),
46            nn.BatchNorm2d(128),
47            nn.ReLU(inplace=True),
48            nn.Conv2d(128, 128, kernel_size=3, padding=1),
49            nn.BatchNorm2d(128),
50            nn.ReLU(inplace=True),
51            nn.MaxPool2d(2, 2),
52            nn.Dropout(0.25),
53            # Block 3
54            nn.Conv2d(128, 256, kernel_size=3, padding=1),
55            nn.BatchNorm2d(256),
56            nn.ReLU(inplace=True),
57            nn.Conv2d(256, 256, kernel_size=3, padding=1),
58            nn.BatchNorm2d(256),
59            nn.ReLU(inplace=True),
60            nn.AdaptiveAvgPool2d((4, 6)),
61        )
62
63        self.classifier = nn.Sequential(
64            nn.Linear(256 * 4 * 6, 512),
65            nn.BatchNorm1d(512),
66            nn.ReLU(inplace=True),
67            nn.Dropout(0.5),
68            nn.Linear(512, 256),
69            nn.BatchNorm1d(256),
70            nn.ReLU(inplace=True),
71            nn.Dropout(0.3),
72            nn.Linear(256, num_classes),
73        )

Initializes the CNN3 model architecture.

Arguments:
  • num_classes (int): Number of target classes.
features
classifier
def forward(self, x): 
75    def forward(self, x):
76        """
77        Defines the forward pass of the CNN3 model.
78
79        Args:
80            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
81
82        Returns:
83            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
84        """
85        x = self.features(x)
86        x = x.view(x.size(0), -1)
87        x = self.classifier(x)
88        return x

Defines the forward pass of the CNN3 model.

Arguments:
  • x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
Returns:

torch.Tensor: Output logits tensor of shape (batch_size, num_classes).

class CNN4(torch.nn.modules.module.Module): 
 91class CNN4(nn.Module):
 92    """
 93    CNN4 is a deeper convolutional neural network with four convolutional blocks.
 94
 95    This architecture is designed for improved feature extraction by adding an extra conv block
 96    and adaptive pooling before classification layers.
 97
 98    Args:
 99        num_classes (int): Number of output classes for classification.
100
101    Attributes:
102        features (nn.Sequential): Convolutional feature extractor with four blocks.
103        classifier (nn.Sequential): Fully connected classifier.
104    """
105
106    def __init__(self, num_classes):
107        """
108        Initializes the CNN4 model architecture.
109
110        Args:
111            num_classes (int): Number of target classes.
112        """
113        super(CNN4, self).__init__()
114        self.features = nn.Sequential(
115            # Block 1
116            nn.Conv2d(3, 64, kernel_size=3, padding=1),
117            nn.BatchNorm2d(64),
118            nn.ReLU(inplace=True),
119            nn.Conv2d(64, 64, kernel_size=3, padding=1),
120            nn.BatchNorm2d(64),
121            nn.ReLU(inplace=True),
122            nn.MaxPool2d(2, 2),
123            nn.Dropout(0.25),
124            # Block 2
125            nn.Conv2d(64, 128, kernel_size=3, padding=1),
126            nn.BatchNorm2d(128),
127            nn.ReLU(inplace=True),
128            nn.Conv2d(128, 128, kernel_size=3, padding=1),
129            nn.BatchNorm2d(128),
130            nn.ReLU(inplace=True),
131            nn.MaxPool2d(2, 2),
132            nn.Dropout(0.25),
133            # Block 3
134            nn.Conv2d(128, 256, kernel_size=3, padding=1),
135            nn.BatchNorm2d(256),
136            nn.ReLU(inplace=True),
137            nn.Conv2d(256, 256, kernel_size=3, padding=1),
138            nn.BatchNorm2d(256),
139            nn.ReLU(inplace=True),
140            nn.MaxPool2d(2, 2),
141            nn.Dropout(0.25),
142            # Block 4 (Additional block for more depth)
143            nn.Conv2d(256, 512, kernel_size=3, padding=1),
144            nn.BatchNorm2d(512),
145            nn.ReLU(inplace=True),
146            nn.Conv2d(512, 512, kernel_size=3, padding=1),
147            nn.BatchNorm2d(512),
148            nn.ReLU(inplace=True),
149            # Replace MaxPool with adaptive pooling to fix feature map size
150            nn.AdaptiveAvgPool2d((4, 6)),
151        )
152
153        self.classifier = nn.Sequential(
154            nn.Linear(512 * 4 * 6, 1024),
155            nn.BatchNorm1d(1024),
156            nn.ReLU(inplace=True),
157            nn.Dropout(0.5),
158            nn.Linear(1024, 512),
159            nn.BatchNorm1d(512),
160            nn.ReLU(inplace=True),
161            nn.Dropout(0.3),
162            nn.Linear(512, num_classes),
163        )
164
165    def forward(self, x):
166        """
167        Defines the forward pass of the CNN4 model.
168
169        Args:
170            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
171
172        Returns:
173            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
174        """
175        x = self.features(x)
176        x = x.view(x.size(0), -1)
177        x = self.classifier(x)
178        return x

CNN4 is a deeper convolutional neural network with four convolutional blocks.

This architecture is designed for improved feature extraction by adding an extra conv block and adaptive pooling before classification layers.

Arguments:
  • num_classes (int): Number of output classes for classification.
Attributes:
  • features (nn.Sequential): Convolutional feature extractor with four blocks.
  • classifier (nn.Sequential): Fully connected classifier.
CNN4(num_classes) 
106    def __init__(self, num_classes):
107        """
108        Initializes the CNN4 model architecture.
109
110        Args:
111            num_classes (int): Number of target classes.
112        """
113        super(CNN4, self).__init__()
114        self.features = nn.Sequential(
115            # Block 1
116            nn.Conv2d(3, 64, kernel_size=3, padding=1),
117            nn.BatchNorm2d(64),
118            nn.ReLU(inplace=True),
119            nn.Conv2d(64, 64, kernel_size=3, padding=1),
120            nn.BatchNorm2d(64),
121            nn.ReLU(inplace=True),
122            nn.MaxPool2d(2, 2),
123            nn.Dropout(0.25),
124            # Block 2
125            nn.Conv2d(64, 128, kernel_size=3, padding=1),
126            nn.BatchNorm2d(128),
127            nn.ReLU(inplace=True),
128            nn.Conv2d(128, 128, kernel_size=3, padding=1),
129            nn.BatchNorm2d(128),
130            nn.ReLU(inplace=True),
131            nn.MaxPool2d(2, 2),
132            nn.Dropout(0.25),
133            # Block 3
134            nn.Conv2d(128, 256, kernel_size=3, padding=1),
135            nn.BatchNorm2d(256),
136            nn.ReLU(inplace=True),
137            nn.Conv2d(256, 256, kernel_size=3, padding=1),
138            nn.BatchNorm2d(256),
139            nn.ReLU(inplace=True),
140            nn.MaxPool2d(2, 2),
141            nn.Dropout(0.25),
142            # Block 4 (Additional block for more depth)
143            nn.Conv2d(256, 512, kernel_size=3, padding=1),
144            nn.BatchNorm2d(512),
145            nn.ReLU(inplace=True),
146            nn.Conv2d(512, 512, kernel_size=3, padding=1),
147            nn.BatchNorm2d(512),
148            nn.ReLU(inplace=True),
149            # Replace MaxPool with adaptive pooling to fix feature map size
150            nn.AdaptiveAvgPool2d((4, 6)),
151        )
152
153        self.classifier = nn.Sequential(
154            nn.Linear(512 * 4 * 6, 1024),
155            nn.BatchNorm1d(1024),
156            nn.ReLU(inplace=True),
157            nn.Dropout(0.5),
158            nn.Linear(1024, 512),
159            nn.BatchNorm1d(512),
160            nn.ReLU(inplace=True),
161            nn.Dropout(0.3),
162            nn.Linear(512, num_classes),
163        )

Initializes the CNN4 model architecture.

Arguments:
  • num_classes (int): Number of target classes.
features
classifier
def forward(self, x): 
165    def forward(self, x):
166        """
167        Defines the forward pass of the CNN4 model.
168
169        Args:
170            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
171
172        Returns:
173            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
174        """
175        x = self.features(x)
176        x = x.view(x.size(0), -1)
177        x = self.classifier(x)
178        return x

Defines the forward pass of the CNN4 model.

Arguments:
  • x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
Returns:

torch.Tensor: Output logits tensor of shape (batch_size, num_classes).

class tinyViT(torch.nn.modules.module.Module): 
181class tinyViT(nn.Module):
182    """
183    tinyViT is a vision transformer model based on the DeiT tiny architecture, pretrained on ImageNet.
184
185    It replaces the original classification head with a linear layer matching the number of classes.
186
187    Args:
188        num_classes (int): Number of output classes.
189        img_size (int): Input image size (assumed square).
190        patch_size (int): Size of the image patches used by the transformer.
191
192    Attributes:
193        backbone (nn.Module): The underlying vision transformer model with modified head.
194    """
195
196    def __init__(
197        self, 
198        num_classes: int, 
199        img_size: int, 
200        patch_size: int, 
201        pretrained: bool = False,
202    ):
203        """
204        Initializes the tinyViT model with a pretrained DeiT backbone.
205
206        Args:
207            num_classes (int): Number of target classes.
208            img_size (int): Input image size (height and width).
209            patch_size (int): Patch size for the transformer.
210            pretrained (bool): Whether to use the pretrained model, deault to `False`.
211        """
212        super(tinyViT, self).__init__()
213        self.backbone = timm.create_model(
214            "deit_tiny_patch16_224",
215            pretrained=pretrained,
216            img_size=img_size,
217            patch_size=patch_size,
218        )
219        in_features = self.backbone.head.in_features
220        self.backbone.head = nn.Linear(in_features, num_classes)
221
222    def forward(self, x):
223        """
224        Defines the forward pass of the tinyViT model.
225
226        Args:
227            x (torch.Tensor): Input tensor of shape (batch_size, 3, img_size, img_size).
228
229        Returns:
230            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
231        """
232        return self.backbone(x)

tinyViT is a vision transformer model based on the DeiT tiny architecture, pretrained on ImageNet.

It replaces the original classification head with a linear layer matching the number of classes.

Arguments:
  • num_classes (int): Number of output classes.
  • img_size (int): Input image size (assumed square).
  • patch_size (int): Size of the image patches used by the transformer.
Attributes:
  • backbone (nn.Module): The underlying vision transformer model with modified head.
tinyViT( num_classes: int, img_size: int, patch_size: int, pretrained: bool = False) 
196    def __init__(
197        self, 
198        num_classes: int, 
199        img_size: int, 
200        patch_size: int, 
201        pretrained: bool = False,
202    ):
203        """
204        Initializes the tinyViT model with a pretrained DeiT backbone.
205
206        Args:
207            num_classes (int): Number of target classes.
208            img_size (int): Input image size (height and width).
209            patch_size (int): Patch size for the transformer.
210            pretrained (bool): Whether to use the pretrained model, deault to `False`.
211        """
212        super(tinyViT, self).__init__()
213        self.backbone = timm.create_model(
214            "deit_tiny_patch16_224",
215            pretrained=pretrained,
216            img_size=img_size,
217            patch_size=patch_size,
218        )
219        in_features = self.backbone.head.in_features
220        self.backbone.head = nn.Linear(in_features, num_classes)

Initializes the tinyViT model with a pretrained DeiT backbone.

Arguments:
  • num_classes (int): Number of target classes.
  • img_size (int): Input image size (height and width).
  • patch_size (int): Patch size for the transformer.
  • pretrained (bool): Whether to use the pretrained model, deault to False.
backbone
def forward(self, x): 
222    def forward(self, x):
223        """
224        Defines the forward pass of the tinyViT model.
225
226        Args:
227            x (torch.Tensor): Input tensor of shape (batch_size, 3, img_size, img_size).
228
229        Returns:
230            torch.Tensor: Output logits tensor of shape (batch_size, num_classes).
231        """
232        return self.backbone(x)

Defines the forward pass of the tinyViT model.

Arguments:
  • x (torch.Tensor): Input tensor of shape (batch_size, 3, img_size, img_size).
Returns:

torch.Tensor: Output logits tensor of shape (batch_size, num_classes).