Machine Learning

TensorFlow

CSV Prediction

Data Science

Deep Learning

Training TensorFlow for Predicting a Column in a csv file

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Deep Learning TensorFlow & Keras Machine Learning Data Science

Introduction

Predicting a column in a CSV file is one of the most common introductory machine learning tasks. Given a dataset with multiple columns, you train a model to predict one target column from the other feature columns. TensorFlow and Keras provide a straightforward pipeline for this: load the CSV with pandas, preprocess the data, build a model, train, and evaluate. This guide covers both regression (predicting continuous values) and classification (predicting categories).

Step 1: Load and Inspect the Data

python

1import pandas as pd
2import numpy as np
3
4df = pd.read_csv('data.csv')
5print(df.shape)          # (rows, columns)
6print(df.head())
7print(df.describe())     # Statistics for numeric columns
8print(df.isnull().sum()) # Check for missing values

Step 2: Prepare Features and Target

python

1# Regression example: predict 'price' from other columns
2target_column = 'price'
3
4# Separate features and target
5X = df.drop(columns=[target_column])
6y = df[target_column]
7
8# Handle missing values
9X = X.fillna(X.median(numeric_only=True))
10
11# Encode categorical columns
12X = pd.get_dummies(X, drop_first=True)
13
14print(f"Features: {X.shape[1]} columns")
15print(f"Target: {target_column}")

Step 3: Split into Train/Test Sets

python

1from sklearn.model_selection import train_test_split
2
3X_train, X_test, y_train, y_test = train_test_split(
4    X, y, test_size=0.2, random_state=42
5)
6
7print(f"Train: {X_train.shape[0]} samples")
8print(f"Test:  {X_test.shape[0]} samples")

Step 4: Normalize Features

Neural networks train better when features are on similar scales:

python

1from sklearn.preprocessing import StandardScaler
2
3scaler = StandardScaler()
4X_train_scaled = scaler.fit_transform(X_train)
5X_test_scaled = scaler.transform(X_test)  # Use train statistics

Step 5: Build the Model

Regression Model

python

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(128, activation='relu', input_shape=(X_train_scaled.shape[1],)),
5    tf.keras.layers.Dropout(0.2),
6    tf.keras.layers.Dense(64, activation='relu'),
7    tf.keras.layers.Dropout(0.2),
8    tf.keras.layers.Dense(32, activation='relu'),
9    tf.keras.layers.Dense(1)  # Single output for regression
10])
11
12model.compile(
13    optimizer='adam',
14    loss='mse',
15    metrics=['mae']
16)
17
18model.summary()

Classification Model

python

1# For binary classification (e.g., predict 'churn': yes/no)
2num_classes = y.nunique()
3
4if num_classes == 2:
5    model = tf.keras.Sequential([
6        tf.keras.layers.Dense(128, activation='relu', input_shape=(X_train_scaled.shape[1],)),
7        tf.keras.layers.Dropout(0.3),
8        tf.keras.layers.Dense(64, activation='relu'),
9        tf.keras.layers.Dense(1, activation='sigmoid')
10    ])
11    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
12else:
13    # Multi-class classification
14    model = tf.keras.Sequential([
15        tf.keras.layers.Dense(128, activation='relu', input_shape=(X_train_scaled.shape[1],)),
16        tf.keras.layers.Dropout(0.3),
17        tf.keras.layers.Dense(64, activation='relu'),
18        tf.keras.layers.Dense(num_classes, activation='softmax')
19    ])
20    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

Step 6: Train the Model

python

1history = model.fit(
2    X_train_scaled, y_train,
3    epochs=100,
4    batch_size=32,
5    validation_split=0.2,
6    callbacks=[
7        tf.keras.callbacks.EarlyStopping(
8            monitor='val_loss',
9            patience=10,
10            restore_best_weights=True
11        )
12    ],
13    verbose=1
14)

Step 7: Evaluate and Visualize

python

1import matplotlib.pyplot as plt
2
3# Evaluate on test set
4test_loss, test_metric = model.evaluate(X_test_scaled, y_test)
5print(f"Test Loss: {test_loss:.4f}")
6print(f"Test MAE:  {test_metric:.4f}")  # or Accuracy for classification
7
8# Plot training history
9fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
10
11ax1.plot(history.history['loss'], label='Train')
12ax1.plot(history.history['val_loss'], label='Validation')
13ax1.set_title('Loss')
14ax1.legend()
15
16ax2.plot(history.history['mae'], label='Train')
17ax2.plot(history.history['val_mae'], label='Validation')
18ax2.set_title('MAE')
19ax2.legend()
20
21plt.tight_layout()
22plt.show()

Step 8: Make Predictions

python

1# Predict on test data
2predictions = model.predict(X_test_scaled).flatten()
3
4# For regression: compare predictions vs actual
5for pred, actual in list(zip(predictions, y_test))[:5]:
6    print(f"Predicted: {pred:.2f}, Actual: {actual:.2f}")
7
8# For classification: convert probabilities to labels
9if num_classes == 2:
10    predicted_labels = (predictions > 0.5).astype(int)
11else:
12    predicted_labels = np.argmax(model.predict(X_test_scaled), axis=1)

Using tf.data for Large CSV Files

For CSV files too large to fit in memory, use tf.data:

python

1# Create a tf.data dataset directly from CSV
2dataset = tf.data.experimental.make_csv_dataset(
3    'large_data.csv',
4    batch_size=32,
5    label_name=target_column,
6    num_epochs=1,
7    shuffle=True
8)
9
10# Build a feature layer for mixed types
11feature_columns = []
12for col in X.columns:
13    if X[col].dtype == 'object':
14        vocab = X[col].unique()
15        feature_columns.append(
16            tf.feature_column.indicator_column(
17                tf.feature_column.categorical_column_with_vocabulary_list(col, vocab)
18            ))
19    else:
20        feature_columns.append(tf.feature_column.numeric_column(col))
21
22feature_layer = tf.keras.layers.DenseFeatures(feature_columns)

Complete Minimal Example

python

1import pandas as pd
2import tensorflow as tf
3from sklearn.model_selection import train_test_split
4from sklearn.preprocessing import StandardScaler
5
6# Load
7df = pd.read_csv('housing.csv')
8X = pd.get_dummies(df.drop(columns=['price']), drop_first=True).fillna(0)
9y = df['price']
10
11# Split and scale
12X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
13scaler = StandardScaler()
14X_train = scaler.fit_transform(X_train)
15X_test = scaler.transform(X_test)
16
17# Build, train, evaluate
18model = tf.keras.Sequential([
19    tf.keras.layers.Dense(64, activation='relu', input_shape=(X_train.shape[1],)),
20    tf.keras.layers.Dense(32, activation='relu'),
21    tf.keras.layers.Dense(1)
22])
23model.compile(optimizer='adam', loss='mse', metrics=['mae'])
24model.fit(X_train, y_train, epochs=50, validation_split=0.2, verbose=0)
25loss, mae = model.evaluate(X_test, y_test)
26print(f"Test MAE: {mae:.2f}")

Common Pitfalls

Not scaling features: Neural networks are sensitive to feature scales. A feature ranging from 0-1000 will dominate one ranging from 0-1. Always use StandardScaler or MinMaxScaler before training.
Data leakage: Fitting the scaler on the full dataset (including test data) leaks test information into training. Always fit_transform on training data only, then transform test data.
Forgetting to encode categoricals: Passing string columns directly to a neural network fails. Use pd.get_dummies(), LabelEncoder, or TensorFlow feature columns.
Too many epochs without early stopping: Without EarlyStopping, the model overfits — training loss keeps decreasing while validation loss increases. Always use early stopping with restore_best_weights=True.
Wrong loss function: Use mse for regression, binary_crossentropy for binary classification, and sparse_categorical_crossentropy for multi-class. Using the wrong loss function produces garbage predictions.

Summary

Load CSV with pandas, separate features (X) and target (y)
Encode categoricals with pd.get_dummies(), fill missing values, scale features with StandardScaler
Use train_test_split to create train/test sets — fit the scaler on train only
Build a Sequential model with Dense layers, Dropout for regularization, and the correct output activation
Train with EarlyStopping and validation_split to prevent overfitting
Evaluate with model.evaluate() and inspect the training history plot for convergence

System design interview problems on this topic

All system design practice problems

Free course

Beginner

7 lessons

2 hours

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Basis Lab

No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Training TensorFlow for Predicting a Column in a csv file

ML System Design practice on Codemia

Introduction

Step 1: Load and Inspect the Data

Step 2: Prepare Features and Target

Step 3: Split into Train/Test Sets

Step 4: Normalize Features

Step 5: Build the Model

Regression Model

Classification Model

Step 6: Train the Model

Step 7: Evaluate and Visualize

Step 8: Make Predictions

Using tf.data for Large CSV Files

Complete Minimal Example

Common Pitfalls

Summary

Related reading

System design interview problems on this topic

Tackling System Design Interview Problems

No course covers this one yet

ML System Design practice on Codemia