import torch
from torchvision import models

# weights=None: inspect the architecture without downloading weights.
model = models.resnet18(weights=None).eval()
x = torch.randn(1, 3, 224, 224)

def inspect(name):
    def hook(module, inputs, output):
        print(name, tuple(output.shape))
    return hook

handles = [getattr(model, name).register_forward_hook(inspect(name))
           for name in ["maxpool", "layer1", "layer2", "layer3", "layer4", "avgpool", "fc"]]
with torch.inference_mode():
    logits = model(x)
for handle in handles:
    handle.remove()
print("Total parameters:", sum(p.numel() for p in model.parameters()))
# A basic identity block returns ReLU(F(x) + x).
# A shape-changing block uses ReLU(F(x) + projection(x)).
# For ResNet-50, torchvision uses bottlenecks (1x1, 3x3, 1x1).
