| import torch |
| from sklearn.datasets import load_iris |
| from torch.utils.data import DataLoader, Dataset |
| import torch.nn as nn |
| import torch.nn.functional as fnn |
| import torch.optim as optim |
|
|
| class IrisDataset(Dataset): |
| def __init__(self): |
| data = load_iris() |
| self.data = torch.tensor(data.data, dtype=torch.float32) |
| self.targets = torch.tensor(data.target, dtype=torch.long) |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, idx): |
| return self.data[idx], self.targets[idx] |
|
|
| dataset = IrisDataset() |
| dataloader = DataLoader(dataset, batch_size=16, shuffle=True) |
|
|
| class INet(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.fc1 = nn.Linear(4, 64) |
| self.fc2 = nn.Linear(64, 32) |
| self.fc3 = nn.Linear(32, 16) |
| self.fc4 = nn.Linear(16, 8) |
| self.fc5 = nn.Linear(8, 4) |
|
|
| def forward(self, x): |
| x = self.fc1(x) |
| x = fnn.relu(x) |
| x = self.fc2(x) |
| x = fnn.relu(x) |
| x = self.fc3(x) |
| x = fnn.relu(x) |
| x = self.fc4(x) |
| x = fnn.relu(x) |
| x = self.fc5(x) |
| return x |
|
|
| model = INet() |
| criterion = nn.CrossEntropyLoss() |
| optimizer = optim.Adam(model.parameters(), lr=0.01) |
|
|
| for epoch in range(30): |
| for X, y in dataloader: |
| preds = model(X) |
| loss = criterion(preds, y) |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
|
|
| print(f"Epoch {epoch}, Loss: {loss.item()}") |
|
|
| torch.save(model.state_dict(), "/models/inet.pth") |
|
|
