File size: 1,154 Bytes
bf26a42 |
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 |
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from models.moe_model import MoEModel
from utils.data_loader import load_data
# Load data
train_loader, test_loader = load_data()
# Initialize model, loss function, and optimizer
model = MoEModel(input_dim=512, num_experts=3)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Training loop
for epoch in range(10):
model.train()
for vision_input, audio_input, sensor_input, labels in train_loader:
optimizer.zero_grad()
outputs = model(vision_input, audio_input, sensor_input)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}, Loss: {loss.item()}")
# Evaluation
model.eval()
correct, total = 0, 0
with torch.no_grad():
for vision_input, audio_input, sensor_input, labels in test_loader:
outputs = model(vision_input, audio_input, sensor_input)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Accuracy: {100 * correct / total}%")
|