Spaces:
Sleeping
Sleeping
File size: 2,262 Bytes
a18e884 | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvBlock(nn.Module):
def __init__(self, in_channels: int, out_channels: int, dropout_rate: float = 0.25):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels,kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Dropout2d(p=dropout_rate),
)
def forward(self, x):
return self.block(x)
class SaraCNN(nn.Module):
def __init__(self, num_classes: int = 6):
super().__init__()
self.features = nn.Sequential(
ConvBlock(3, 32, dropout_rate=0.25),
ConvBlock(32, 64, dropout_rate=0.25),
ConvBlock(64, 128, dropout_rate=0.25),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 18 * 18, 256),
nn.BatchNorm1d(256),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(256, num_classes),
)
self._init_weights()
def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
elif isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.features(x)
x = self.classifier(x)
return x
if __name__ == "__main__":
model = SaraCNN(num_classes=6)
dummy = torch.randn(4, 3, 150, 150)
out = model(dummy)
print("SaraCNN output shape:", out.shape)
total = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable parameters: {total:,}")
|