Views
No views yet
1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class Model(nn.Module):
6 def __init__(self,in_features=4,h1=8,h2=9,out_feauteres=3):
7 super().__init__() # instantiate our nn.Module
8 self.fc1 = nn.Linear(in_features,h1)
9 self.fc2 = nn.Linear(h1,h2)
10 self.out = nn.Linear(h2,out_feauteres)
11
12 def forward(self,x):
13 x = F.relu(self.fc1(x))
14 x = F.relu(self.fc2(x))
15 x = self.out(x)
16 return x
17
18def number_to_follower(x):
19 if x == 0:
20 return'Setosa'
21 elif x == 1:
22 return 'Versicolor'
23 elif x == 2:
24 return 'Virginica'
25
26model = Model()
27
28model.load_state_dict(torch.load('ombayus_iris_model.pt'))
29
30new_iris = torch.tensor([5.9,3.0,5.1,1.8])
31
32with torch.no_grad():
33 print(number_to_follower(model(new_iris).argmax().item()))
34