34 lines
942 B
Python
34 lines
942 B
Python
import torch
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
|
|
file = './RegressionModels/PredictSallary/Salary_dataset.csv'
|
|
df = pd.read_csv(file)
|
|
|
|
X = torch.tensor(df["YearsExperience"].values, dtype=torch.float32).unsqueeze(1)
|
|
y = torch.tensor(df['Salary'].values, dtype=torch.float32).unsqueeze(1)
|
|
|
|
model = torch.nn.Sequential(
|
|
torch.nn.Linear(1,1)
|
|
)
|
|
|
|
loss_fn = torch.nn.MSELoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
|
|
|
|
# 5. Training loop
|
|
for epoch in range(1000):
|
|
y_pred = model(X)
|
|
loss = loss_fn(y_pred, y)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
if epoch % 100 == 99:
|
|
print(f'Epoch {epoch+1}, Loss: {loss.item():.2f}')
|
|
|
|
# 6. Plot the results
|
|
plt.scatter(X.numpy(), y.numpy(), label='Actual data')
|
|
plt.plot(X.numpy(), model(X).detach().numpy(), color='red', label='Model prediction')
|
|
plt.xlabel('Years of Experience')
|
|
plt.ylabel('Salary')
|
|
plt.legend()
|
|
plt.show() |