Spaces:
Runtime error
Runtime error
File size: 883 Bytes
710f946 de372c9 710f946 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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)
|