nn.Module has two Keras counterparts: the Sequential API for simple stacks and the Functional API for anything with branches. Map your PyTorch model-building instincts across.
Why: keras.Sequential is the direct analog of nn.Sequential — a linear stack of layers — but with one big convenience: Keras infers input shapes, so you do not specify in_features on every layer like you do in PyTorch. When: use it for straightforward feed-forward or CNN stacks. Where: keras.layers.Dense is nn.Linear, and the activation is usually a keyword argument rather than a separate layer.
from tensorflow import keras
from tensorflow.keras import layers
# PyTorch: nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 3))
model = keras.Sequential([
layers.Dense(64, activation="relu", input_shape=(10,)), # = nn.Linear + ReLU
layers.Dense(3), # output logits
])
model.summary() # like printing an nn.Module, but with shapes + param countsWhy: when a model is not a simple stack — skip connections, multiple inputs/outputs, shared layers — the Functional API builds it by calling layers on tensors, much like writing an nn.Module’s forward method, but declaratively. When: reach for it for anything a Sequential cannot express. Where: it is Keras’s equivalent of a custom forward pass, and most real architectures use it.
from tensorflow import keras
from tensorflow.keras import layers
inputs = keras.Input(shape=(10,))
x = layers.Dense(64, activation="relu")(inputs) # call layers on tensors
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(3)(x)
model = keras.Model(inputs, outputs) # like defining forward(), declaratively
# Skip connections, multi-input/output, shared layers all work here.Why: if you want the exact feel of PyTorch — define layers in __init__ and the forward pass in a method — Keras supports model subclassing, where you override call() instead of forward(). When: use it for dynamic or research models where you need full imperative control. Where: this is the most PyTorch-like option, and a good fallback when the Functional API feels constraining.
from tensorflow import keras
from tensorflow.keras import layers
class MLP(keras.Model):
def __init__(self, hidden, out):
super().__init__()
self.d1 = layers.Dense(hidden, activation="relu") # like __init__ layers
self.d2 = layers.Dense(out)
def call(self, x): # 'call' is Keras's 'forward'
return self.d2(self.d1(x))
model = MLP(hidden=64, out=3)