54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
|
import os
|
||
|
import pandas as pd
|
||
|
import torch
|
||
|
import torch.nn as nn
|
||
|
import torch.optim as optim
|
||
|
|
||
|
file_path = os.path.join("C:", os.sep, "Users", "reyva", "OneDrive", "Pulpit", "studia", "InzynieriaUczeniaMaszynowego", 'Train1.csv')
|
||
|
|
||
|
train_data = pd.read_csv(file_path)
|
||
|
|
||
|
train_data = train_data.dropna(subset=['price_range'])
|
||
|
|
||
|
valid_values = {0.0, 1.0, 2.0, 3.0}
|
||
|
assert set(train_data['price_range'].unique()) <= valid_values, "Unexpected values in price_range"
|
||
|
|
||
|
class SimpleNN(nn.Module):
|
||
|
def __init__(self, input_size, hidden_size, output_size):
|
||
|
super(SimpleNN, self).__init__()
|
||
|
self.fc1 = nn.Linear(input_size, hidden_size)
|
||
|
self.fc2 = nn.Linear(hidden_size, output_size)
|
||
|
self.softmax = nn.Softmax(dim=1)
|
||
|
|
||
|
def forward(self, x):
|
||
|
x = torch.relu(self.fc1(x))
|
||
|
x = self.fc2(x)
|
||
|
return self.softmax(x)
|
||
|
|
||
|
input_size = len(train_data.columns) - 2
|
||
|
hidden_size = 50
|
||
|
output_size = len(valid_values)
|
||
|
|
||
|
model = SimpleNN(input_size, hidden_size, output_size)
|
||
|
|
||
|
criterion = nn.CrossEntropyLoss()
|
||
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
||
|
|
||
|
epochs = 10
|
||
|
for epoch in range(epochs):
|
||
|
inputs = torch.tensor(train_data.drop(['price_range', 'ID'], axis=1).values, dtype=torch.float32)
|
||
|
labels = torch.tensor(train_data['price_range'].values, dtype=torch.long)
|
||
|
|
||
|
optimizer.zero_grad()
|
||
|
|
||
|
outputs = model(inputs)
|
||
|
loss = criterion(outputs, labels)
|
||
|
loss.backward()
|
||
|
optimizer.step()
|
||
|
|
||
|
print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item()}")
|
||
|
|
||
|
save_path = os.path.join("C:", os.sep, "Users", "reyva", "OneDrive", "Pulpit", "studia", "InzynieriaUczeniaMaszynowego", "model.pth")
|
||
|
torch.save(model.state_dict(), save_path)
|
||
|
|