File size: 1,532 Bytes
bfd7e63 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | 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")
|