Spaces:
Sleeping
Sleeping
Delete tasks/Model_Loader.py
Browse files- tasks/Model_Loader.py +0 -42
tasks/Model_Loader.py
DELETED
|
@@ -1,42 +0,0 @@
|
|
| 1 |
-
import torch
|
| 2 |
-
|
| 3 |
-
class M5(torch.nn.Module):
|
| 4 |
-
def __init__(self, num_classes=2): # Ensure it matches dataset labels (chainsaw/environment)
|
| 5 |
-
super(M5, self).__init__()
|
| 6 |
-
self.conv1 = torch.nn.Conv1d(in_channels=1, out_channels=32, kernel_size=80, stride=4)
|
| 7 |
-
self.bn1 = torch.nn.BatchNorm1d(32)
|
| 8 |
-
self.conv2 = torch.nn.Conv1d(in_channels=32, out_channels=64, kernel_size=3)
|
| 9 |
-
self.bn2 = torch.nn.BatchNorm1d(64)
|
| 10 |
-
self.conv3 = torch.nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3)
|
| 11 |
-
self.bn3 = torch.nn.BatchNorm1d(128)
|
| 12 |
-
self.conv4 = torch.nn.Conv1d(in_channels=128, out_channels=256, kernel_size=3)
|
| 13 |
-
self.bn4 = torch.nn.BatchNorm1d(256)
|
| 14 |
-
self.fc1 = torch.nn.Linear(256, num_classes)
|
| 15 |
-
|
| 16 |
-
def forward(self, x):
|
| 17 |
-
x = torch.nn.functional.relu(self.bn1(self.conv1(x)))
|
| 18 |
-
x = torch.nn.functional.max_pool1d(x, 4)
|
| 19 |
-
x = torch.nn.functional.relu(self.bn2(self.conv2(x)))
|
| 20 |
-
x = torch.nn.functional.max_pool1d(x, 4)
|
| 21 |
-
x = torch.nn.functional.relu(self.bn3(self.conv3(x)))
|
| 22 |
-
x = torch.nn.functional.max_pool1d(x, 4)
|
| 23 |
-
x = torch.nn.functional.relu(self.bn4(self.conv4(x)))
|
| 24 |
-
x = torch.nn.functional.max_pool1d(x, 4)
|
| 25 |
-
x = torch.mean(x, dim=2)
|
| 26 |
-
x = self.fc1(x)
|
| 27 |
-
return x
|
| 28 |
-
|
| 29 |
-
def load_model(model_path, num_classes=2):
|
| 30 |
-
"""
|
| 31 |
-
Load trained M5 model.
|
| 32 |
-
"""
|
| 33 |
-
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 34 |
-
model = M5(num_classes=num_classes).to(device)
|
| 35 |
-
model.load_state_dict(torch.load(model_path, map_location=device))
|
| 36 |
-
model.eval() # Set model to evaluation mode
|
| 37 |
-
return model, device
|
| 38 |
-
|
| 39 |
-
if __name__ == "__main__":
|
| 40 |
-
model, device = load_model("quantized_teacher_m5_static.pth")
|
| 41 |
-
print("✅ Model successfully loaded!")
|
| 42 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|