PyTorch-casas-SA / train.py
natisol87's picture
Upload 8 files
d2e67f5 verified
Raw History Blame Contribute Delete
3.31 kB
"""Entrena una regresi贸n lineal con PyTorch y la exporta a model.safetensors + config.json.
Mismo problema que el ejemplo 01: precio de una casa a partir de m虏, balc贸n y ascensor.
Diferencia importante con scikit-learn: PyTorch guarda solo los pesos (state_dict).
La clase del modelo y la normalizaci贸n de las entradas no viajan en el archivo de pesos,
as铆 que la normalizaci贸n se guarda aparte, en config.json.
Uso:
python train.py
"""
import json
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from safetensors.torch import save_file
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from modelo import RegresionLineal
CARPETA = Path(__file__).parent
ENTRADAS = ["m2", "balcon", "ascensor"]
torch.manual_seed(42)
def generar_casas(n=200, semilla=42):
"""Dataset sint茅tico de la Clase 10 (misma semilla, mismos datos)."""
rng = np.random.default_rng(semilla)
m2 = rng.integers(35, 121, n)
balcon = rng.integers(0, 2, n)
ascensor = rng.integers(0, 2, n)
precio = 1.5 * m2 + 8 * balcon + 12 * ascensor + 20 + rng.normal(0, 4, n)
return pd.DataFrame({"m2": m2, "balcon": balcon, "ascensor": ascensor, "precio": precio.round(1)})
# 1. Datos
casas = generar_casas()
X = casas[ENTRADAS].to_numpy(dtype="float32")
y = casas["precio"].to_numpy(dtype="float32").reshape(-1, 1)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Normalizaci贸n: se calcula SOLO con los datos de entrenamiento y se guarda para predecir
media = X_train.mean(axis=0)
desvio = X_train.std(axis=0)
X_train_t = torch.tensor((X_train - media) / desvio)
y_train_t = torch.tensor(y_train)
# 2. Modelo, funci贸n de p茅rdida y optimizador
modelo = RegresionLineal(n_entradas=len(ENTRADAS))
perdida_fn = torch.nn.MSELoss()
optimizador = torch.optim.SGD(modelo.parameters(), lr=0.1)
# 3. El ciclo de entrenamiento
for epoca in range(500):
prediccion = modelo(X_train_t) # el modelo predice
perdida = perdida_fn(prediccion, y_train_t) # la p茅rdida mide el error
optimizador.zero_grad()
perdida.backward() # se calculan los gradientes
optimizador.step() # el optimizador corrige pesos y bias
if epoca % 100 == 0:
print(f"茅poca {epoca:3d} 路 p茅rdida (MSE) {perdida.item():.2f}")
# 4. Evaluaci贸n con datos que el modelo no vio
modelo.eval()
with torch.no_grad():
pred_test = modelo(torch.tensor((X_test - media) / desvio)).numpy()
print(f"MAE en test: {mean_absolute_error(y_test, pred_test):.2f} miles de US$")
print(f"R虏 en test: {r2_score(y_test, pred_test):.3f}")
# 5. Exportaci贸n: pesos en safetensors, todo lo dem谩s en config.json
save_file(modelo.state_dict(), CARPETA / "model.safetensors")
config = {
"arquitectura": "RegresionLineal (torch.nn.Linear)",
"entradas": ENTRADAS,
"salida": "precio (miles de US$)",
"normalizacion": {"media": media.tolist(), "desvio": desvio.tolist()},
}
(CARPETA / "config.json").write_text(json.dumps(config, indent=2, ensure_ascii=False), encoding="utf-8")
print("Modelo guardado en model.safetensors + config.json")