ium_487197/ium_train.py

65 lines
2.0 KiB
Python
Raw Normal View History

2023-05-11 17:47:16 +02:00
# This is a sample Python script.
# Press Shift+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
2023-05-11 19:17:17 +02:00
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam
2023-05-11 17:47:16 +02:00
import pandas as pd
import tensorflow as tf
import numpy as np
from sklearn.preprocessing import LabelEncoder
2023-05-11 20:13:03 +02:00
import argparse
2023-05-11 17:47:16 +02:00
def get_x_y(data):
lb = LabelEncoder()
data = data.drop(["Location 1"], axis=1)
data = data.drop(columns=["Longitude", "Latitude", "Location", "Total Incidents", "CrimeTime", "Neighborhood", "Post", "CrimeDate", "Inside/Outside"], axis=1)
for column_name in data.columns:
data[column_name] = lb.fit_transform(data[column_name])
x = data.drop('Weapon', axis=1)
y = data['Weapon']
return data, x, y
def train_model():
2023-05-11 20:13:03 +02:00
parser = argparse.ArgumentParser(description='Train')
parser.add_argument('-epochs', type=int, default=20)
parser.add_argument('-lr', type=float, default=0.01)
parser.add_argument('-validation_split', type=int, default=0.2)
args = parser.parse_args()
2023-05-11 17:47:16 +02:00
train = pd.read_csv('baltimore_train.csv')
data_train, x_train, y_train = get_x_y(train)
normalizer = tf.keras.layers.Normalization(axis=1)
normalizer.adapt(np.array(x_train))
model = Sequential(normalizer)
model.add(Dense(64, activation="relu"))
model.add(Dense(10, activation='relu'))
model.add(Dense(10, activation='relu'))
model.add(Dense(10, activation='relu'))
model.add(Dense(5, activation="softmax"))
2023-05-11 20:13:03 +02:00
model.compile(Adam(learning_rate=args.lr), loss='sparse_categorical_crossentropy', metrics = ['accuracy'] )
2023-05-11 17:47:16 +02:00
model.summary()
history = model.fit(
x_train,
y_train,
2023-05-11 20:13:03 +02:00
epochs=args.epochs,
validation_split=args.validation_split)
2023-05-11 17:47:16 +02:00
hist = pd.DataFrame(history.history)
hist['epoch'] = history.epoch
2023-05-11 19:45:43 +02:00
model.save('baltimore_model.h5')
2023-05-11 17:47:16 +02:00
train_model()