seronk commited on
Commit
5f79afa
·
verified ·
1 Parent(s): c8d4026

Delete tasks/Model_Loader.py

Browse files
Files changed (1) hide show
  1. 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
-