Spaces:
Sleeping
Sleeping
File size: 1,799 Bytes
163b2c6 f95c051 163b2c6 f95c051 163b2c6 f95c051 |
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 |
import torch
class M5(torch.nn.Module):
def __init__(self, num_classes=2): # Ensure it matches dataset labels (chainsaw/environment)
super(M5, self).__init__()
self.conv1 = torch.nn.Conv1d(in_channels=1, out_channels=32, kernel_size=80, stride=4)
self.bn1 = torch.nn.BatchNorm1d(32)
self.conv2 = torch.nn.Conv1d(in_channels=32, out_channels=64, kernel_size=3)
self.bn2 = torch.nn.BatchNorm1d(64)
self.conv3 = torch.nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3)
self.bn3 = torch.nn.BatchNorm1d(128)
self.conv4 = torch.nn.Conv1d(in_channels=128, out_channels=256, kernel_size=3)
self.bn4 = torch.nn.BatchNorm1d(256)
self.fc1 = torch.nn.Linear(256, num_classes)
def forward(self, x):
x = torch.nn.functional.relu(self.bn1(self.conv1(x)))
x = torch.nn.functional.max_pool1d(x, 4)
x = torch.nn.functional.relu(self.bn2(self.conv2(x)))
x = torch.nn.functional.max_pool1d(x, 4)
x = torch.nn.functional.relu(self.bn3(self.conv3(x)))
x = torch.nn.functional.max_pool1d(x, 4)
x = torch.nn.functional.relu(self.bn4(self.conv4(x)))
x = torch.nn.functional.max_pool1d(x, 4)
x = torch.mean(x, dim=2)
x = self.fc1(x)
return x
def load_model(model_path, num_classes=2):
"""
Load trained M5 model.
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = M5(num_classes=num_classes).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval() # Set model to evaluation mode
return model, device
if __name__ == "__main__":
model, device = load_model("quantized_teacher_m5_static.pth")
print("✅ Model successfully loaded!")
|