Keras vs PyTorch Code Comparison: From Data Loading to Prediction
If you’re stepping into the world of deep learning, one of the biggest choices you’ll make is between Keras (with TensorFlow backend) and PyTorch. Both are powerful, well-supported frameworks, but they differ in their syntax, flexibility, and learning curve.
In this blog post, I walk you through a full machine learning pipeline using both Keras and PyTorch — from data loading to model prediction — so you can see how they compare step by step.
📦 1. Loading Data
🔹 PyTorch: Custom Dataset with Dataset and DataLoader
class CSVDataset(Dataset):
# reading the csv and defining predictor and output columns
def __init__(self):
# store the input and output features
self.X = df.values[:,:-1]
self.y = df.values[:,-1]
# ensure all data is numerical - type(float)
self.X = self.X.astype('float32')
self.y = self.y.astype('float32')
# number of rows in dataset
def __len__(self):
return len(self.X)
# get a row at an index
def __getitem__(self, index):
return [self.X[index], self.y[index]]
# split into train and testset - using `random_split`
def get_splits(self, split_ratio = 0.2):
test_size = round(split_ratio * len(self.X))
train_size = len(self.X) - test_size
return random_split(self, [train_size, test_size])
🔹 Keras: Load with Pandas
import pandas as pd
df = pd.read_csv('data.csv')
Takeaway: PyTorch gives you full control over the data pipeline using
DatasetandDataLoader, while Keras keeps it simple and clean for quick experimentation.
🧠 2. Defining the Model
🔹 PyTorch: Manual Architecture with Custom Class
class myNeuralNetwork(Module):
def __init__(self, n_inputs):
# calling constructor of parent class
super().__init__()
# defining the inputs to the first hidden layer
self.hid1 = Linear(n_inputs, 8)
kaiming_uniform_(self.hid1.weight, nonlinearity='relu')
self.act1 = ReLU()
# defining the inputs to the second hidden layer
self.hid2 = Linear(8, 16)
kaiming_uniform_(self.hid2.weight, nonlinearity='relu')
self.act2 = ReLU()
# defining the inputs to the third hidden layer
self.hid3 = Linear(16, 2)
xavier_uniform_(self.hid3.weight)
self.act3 = Softmax(dim=1)
def forward(self, X):
#input and act for layer 1
X = self.hid1(X)
X = self.act1(X)
#input and act for layer 2
X = self.hid2(X)
X = self.act2(X)
#input and act for layer 3
X = self.hid3(X)
X = self.act3(X)
return X
🔹 Keras: Sequential API
model = Sequential([
Dense(8, input_shape=(2,), activation='relu'),
Dense(16, activation='relu'),
Dense(2, activation='softmax')
])
Takeaway: PyTorch gives you control at every step, while Keras’s
SequentialAPI is perfect for quick prototyping.
🏗️ 3. Data Preparation for Training
🔹 PyTorch: Use DataLoader
# Preparing the dataset before training the model
# load the dataset
dataset = CSVDataset()
# get the train and test split
train, test = dataset.get_splits()
# prepare dataloaders -
train_dl = DataLoader(train, batch_size = 32, shuffle = True)
test_dl = DataLoader(test, batch_size= 32, shuffle= False)
### 🔹 Keras: Use Scikit-learn’s `train_test_split`
```python
from sklearn.model_selection import train_test_split
train_data, test_data = train_test_split(df, test_size=0.2, random_state=25)
Takeaway: PyTorch’s
DataLoadersupports efficient mini-batching and shuffling out-of-the-box. Keras users often lean onnumpyorscikit-learn.
🚂 4. Training the Model
🔹 PyTorch: Manual Training Loop
# define the network
model = myNeuralNetwork(2) # 2 input features
# define the number of epochs
epochs = 10
# define the optimizer - SGD
optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
# define the loss function
criterion = CrossEntropyLoss()
# iterate through all the epoch
for epoch in range(epochs):
# go through all the batches generated by dataloader
for i, (inputs, targets) in enumerate(train_dl):
# clear the gradients
optimizer.zero_grad()
# compute the model output
yhat = model(inputs)
# calculate loss
loss = criterion(yhat, targets.type(torch.LongTensor))
# credit assignment
loss.backward()
# update model weights
optimizer.step()
🔹 Keras: Built-in .fit() API
# compile model
model.compile(optimizer= Adam(learning_rate = 0.0001),
loss = 'sparse_categorical_crossentropy',
metrics = ['accuracy']
)
# Model Training and Validation
model_m.fit(x = scaled_train_samples_mult,
y = train_labels,
batch_size= 10,
epochs = 30,
validation_split= 0.1,
shuffle = True,
verbose = 2
)
Takeaway: PyTorch requires you to define every training step — great for customization. Keras handles it all for you behind the scenes.
📊 5. Model Evaluation
🔹 PyTorch: Manual Evaluation
# Evaluate the model
predictions, actuals = list(), list()
# loop over all batches in test set
for i, (inputs, targets) in enumerate(test_dl):
# pass input to the model
y_pred = model(inputs)
# retrieve the numpy array
y_pred = y_pred.detach().numpy()
# pick the index of the highest values
res = np.argmax(y_pred, axis = 1)
# actual output
actual = targets.numpy()
actual = actual.reshape(len(actual), 1)
# store the values in respective lists
predictions.append(list(res))
actuals.append(list(actual))
actuals = [val for sublist in vstack(list(chain(*actuals))) for val in sublist]
predictions = [val for sublist in vstack(list(chain(*predictions))) for val in sublist]
🔹 Keras: One-liner Evaluation
model.evaluate(x=Xtest, y=Ytest)
Takeaway: Keras simplifies evaluation with a single command. PyTorch gives you full flexibility if you need custom metrics or logging.
🤖 6. Making Predictions
🔹 PyTorch: Same as Evaluation (manual loop or model(X))
🔹 Keras: Straightforward .predict()
output = model.predict(X)
Takeaway: Prediction is easier and cleaner in Keras, especially for beginners.
🔚 Final Thoughts: Which Should You Choose?
| Feature | PyTorch | Keras |
|---|---|---|
| Flexibility | ✅ Fine-grained control | ❌ Less customizable |
| Ease of Use | ❌ More code needed | ✅ Cleaner syntax |
| Community | ✅ Strong in research | ✅ Strong in production/deployment |
| Debugging | ✅ Pythonic & transparent | ❌ Abstracted |
| Speed of Prototyping | ❌ Slower | ✅ Faster |
✨ Summary
-
Use Keras if you want a fast, beginner-friendly API for deep learning projects.
-
Choose PyTorch if you want total control and are working on research or custom architectures.
🔗 References
👋 About Me
Hi, I’m Shuvangkar Das, a power systems researcher with a Ph.D. in Electrical Engineering from Clarkson University. I work at the intersection of power electronics, DER, IBR, and AI — building greener, smarter, and more stable grids. Currently, I’m a Research Scientist at EPRI (though everything I share here reflects my personal experience, not my employer’s views).
Over the years, I’ve worked on real-world projects involving large scale EMT simulation and firmware development for grid-forming and grid following inverter and reinforcement learning (RL). I also publish technical content and share hands-on insights with the goal of making complex ideas accessible to engineers and researchers.
📺 Subscribe to my YouTube channel, where I share tutorials, code walk-throughs, and research productivity tips.
Connect with me:
Leave a comment