PyTorch = TensorFlow.js (Different Syntax, Same Ideas)
PyTorch nn.Module and TensorFlow.js tf.LayersModel express the same abstractions — layers, forward passes, loss functions, and optimizers — with different syntax.
PyTorch and TensorFlow.js are two dialects of the same language. Dense layers, conv layers, ReLU, Adam optimizer — all the same building blocks. Once you see the mapping, you can read any PyTorch model definition and know exactly what it does.
Most ML research code is written in PyTorch. When you read a paper's reference implementation, a colleague's notebook, or a GitHub repo, it will almost certainly be PyTorch. This lesson gives you the Rosetta Stone: every PyTorch concept mapped to the TF.js equivalent you already know.
Learning Objectives
- ○Read a PyTorch model definition and identify each layer's TF.js equivalent
- ○Map PyTorch's training loop to TF.js model.fit()
- ○Translate PyTorch optimizers, loss functions, and activations to TF.js
- ○Understand nn.Module's forward() pattern vs. TF.js sequential/functional API
Layer Mapping
The layer names are different, but they do the same thing.
Frontend
tf.sequential / tf.model
const model = tf.sequential({ layers: [tf.layers.dense({units: 64})] })Machine Learning
PyTorch nn.Module
self.fc = nn.Linear(input_dim, 64)import torch.nn as nn
# PyTorch layer → TF.js equivalent
nn.Linear(784, 128) # tf.layers.dense({units: 128, inputShape: [784]})
nn.Conv2d(3, 32, 3, padding=1) # tf.layers.conv2d({filters: 32, kernelSize: 3, padding: 'same'})
nn.BatchNorm1d(128) # tf.layers.batchNormalization()
nn.Dropout(0.5) # tf.layers.dropout({rate: 0.5})
nn.ReLU() # tf.layers.activation({activation: 'relu'}) // or inline
nn.Flatten() # tf.layers.flatten()
nn.LSTM(64, 32) # tf.layers.lstm({units: 32})Model Definition: Side by Side
Here's a simple classifier in both frameworks. Read them in parallel — they're the same model.
import torch
import torch.nn as nn
class SignalClassifier(nn.Module):
def __init__(self, input_dim, num_classes):
super().__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, 128),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(128, 64),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(64, num_classes)
)
def forward(self, x):
return self.model(x)
# Usage
model = SignalClassifier(input_dim=784, num_classes=10)
output = model(input_tensor) # calls forward()import * as tf from '@tensorflow/tfjs';
const model = tf.sequential({
layers: [
tf.layers.dense({ units: 128, activation: 'relu', inputShape: [784] }),
tf.layers.dropout({ rate: 0.3 }),
tf.layers.dense({ units: 64, activation: 'relu' }),
tf.layers.dropout({ rate: 0.3 }),
tf.layers.dense({ units: 10 })
]
});
// Usage
const output = model.predict(inputTensor);Same architecture. Same layer sizes. Same activations. Same dropout rates. Different syntax.
Optimizers and Loss Functions
import torch.optim as optim
# PyTorch → TF.js
optim.Adam(model.parameters(), lr=0.001) # tf.train.adam(0.001)
optim.SGD(model.parameters(), lr=0.01) # tf.train.sgd(0.01)
nn.CrossEntropyLoss() # 'categoricalCrossentropy'
nn.MSELoss() # 'meanSquaredError'
nn.BCELoss() # 'binaryCrossentropy'Training Loop Comparison
PyTorch uses an explicit loop. TF.js wraps it in model.fit(). Same steps, different packaging.
# PyTorch training loop
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
for batch_x, batch_y in dataloader:
optimizer.zero_grad() # 1. Clear gradients
output = model(batch_x) # 2. Forward pass
loss = criterion(output, batch_y) # 3. Compute loss
loss.backward() # 4. Backward pass
optimizer.step() # 5. Update weights
print(f"Epoch {epoch}, Loss: {loss.item()}")// TF.js — same 5 steps, wrapped in model.fit()
model.compile({
optimizer: tf.train.adam(0.001),
loss: 'categoricalCrossentropy',
metrics: ['accuracy']
});
await model.fit(trainX, trainY, {
epochs: 10,
batchSize: 32,
callbacks: {
onEpochEnd: (epoch, logs) => {
console.log(`Epoch ${epoch}, Loss: ${logs.loss}`);
}
}
});The five steps in PyTorch's loop — clear gradients, forward pass, compute loss, backward pass, update weights — are exactly what model.fit() does internally on each batch.
Data Loading
from torch.utils.data import DataLoader, TensorDataset
# PyTorch
dataset = TensorDataset(X_tensor, y_tensor)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# TF.js equivalent
# const dataset = tf.data.zip({
# xs: tf.data.array(X),
# ys: tf.data.array(y)
# }).shuffle(1000).batch(32);Challenge
Read a PyTorch model definition and build the equivalent in TensorFlow.js.
Exercise
PyTorch Model to TensorFlow.js
Read the PyTorch model definition below and build the equivalent model in TensorFlow.js. The model is a simple feedforward classifier. Match the layer types, sizes, activations, and dropout rates exactly. Then compile it with the equivalent optimizer and loss function.
Key Takeaways
- ✓PyTorch nn.Module and TF.js tf.sequential express the same architecture with different syntax
- ✓Layer types map 1:1: nn.Linear → dense, nn.Conv2d → conv2d, nn.LSTM → lstm
- ✓PyTorch's explicit training loop does the same 5 steps that model.fit() wraps
- ✓Optimizers and loss functions have direct equivalents in both frameworks