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. The good news: you already know PyTorch. You know it under different names, and the tables below pair each one with the TF.js API you've been using all along.
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 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})The Same Model, Defined Twice
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 . 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 , , compute , , update ) are exactly what model.fit() does internally on each .
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.
Recall Prompt
What are the five steps in a PyTorch training loop, and where do they appear in TF.js?
Lesson Recap
What you learned
- ✓PyTorch nn.Module and tf.sequential express the same architecture with different syntax: nn.Linear maps to tf.layers.dense, nn.Conv2d to tf.layers.conv2d, nn.LSTM to tf.layers.lstm
- ✓PyTorch's explicit training loop runs the same five steps (zero grads, forward, loss, backward, update) that tf.model.fit() wraps automatically
- ✓Optimizers and loss functions have direct equivalents: optim.Adam -> tf.train.adam, nn.CrossEntropyLoss -> 'categoricalCrossentropy'
The bridge
A PyTorch `nn.Sequential` block and a `tf.sequential` model define the same data flow, same layer order, same shapes; only the method names and loop structure differ.
You can now
Read a PyTorch model definition, identify each layer's TF.js equivalent, and understand how its training loop maps to model.fit().