Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| import torchvision.models as models | |
| class DenseNet121(nn.Module): | |
| """DenseNet121 classifier head for the demo app.""" | |
| def __init__(self, num_classes: int, dropout_rate: float = 0.25): | |
| super().__init__() | |
| self.densenet_model = models.densenet121(weights=None) | |
| self.dropout = nn.Dropout(p=dropout_rate) | |
| self.classifier = nn.Linear(1024, num_classes) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if x.dim() == 3: | |
| x = x.unsqueeze(0) | |
| features = self.densenet_model.features(x) | |
| features = F.gelu(features) | |
| pooled = F.adaptive_avg_pool2d(features, (1, 1)) | |
| flattened = torch.flatten(pooled, 1) | |
| if self.training: | |
| flattened = self.dropout(flattened) | |
| return self.classifier(flattened) | |