Computer Vision Engineering: CNNs, ViTs, and Production

Computer vision is one of the most mature ML fields with a clear succession of paradigms: hand-crafted features (HOG, SIFT) → CNNs → Vision Transformers. If you understand why each transition happened, you understand the field.

CNNs: The Foundation

Convolutional Neural Networks exploit the spatial structure of images: nearby pixels are correlated, and the same local patterns (edges, textures) appear throughout an image regardless of location. Convolutions detect these patterns anywhere in the image (translation equivariance) with far fewer parameters than a fully connected network.

pythonCopy
import torch
import torch.nn as nn

class ConvBlock(nn.Module):
    """Conv → BatchNorm → ReLU  -  the standard building block."""
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            ConvBlock(3, 32),        # 3 channels (RGB) → 32 feature maps
            ConvBlock(32, 64),
            nn.MaxPool2d(2, 2),      # downsample 2x
            ConvBlock(64, 128),
            ConvBlock(128, 128),
            nn.MaxPool2d(2, 2),
            ConvBlock(128, 256),
            nn.AdaptiveAvgPool2d(1)  # global average pooling → 256 × 1 × 1
        )
        self.classifier = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.flatten(1)
        return self.classifier(x)

Transfer Learning: The Standard Approach

Training a CNN from scratch requires millions of images. Transfer learning uses pretrained features from ImageNet (1.2M images, 1000 classes) and fine-tunes them on your task:

pythonCopy
import torchvision.models as models
import torch.nn as nn

# Load pretrained ResNet-50
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# Freeze all layers except the final classifier
for param in model.parameters():
    param.requires_grad = False

# Replace the classifier for your number of classes
n_classes = 5  # your task
model.fc = nn.Linear(model.fc.in_features, n_classes)
# Only model.fc parameters will be updated

# Training: use small learning rate for fine-tuning
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

# Full fine-tuning (unfreeze all): even smaller learning rate
for param in model.parameters():
    param.requires_grad = True
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

When to freeze vs. fine-tune all layers:

• Freeze all: your dataset is small (<1000 images) and similar to ImageNet
• Fine-tune all: your dataset is large or domain-shifted from ImageNet (medical images, satellite imagery)

Data Augmentation: Your Best Regularizer

Augmentation artificially expands your dataset and teaches invariances:

pythonCopy
from torchvision import transforms

# Training transforms (with augmentation)
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1),
    transforms.RandomGrayscale(p=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],  # ImageNet statistics
                         std=[0.229, 0.224, 0.225])
])

# Validation transforms (no augmentation  -  evaluate on clean images)
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

Choose augmentations that are plausible for your domain. For medical images: rotations and flips yes, extreme color jitter no. For satellite imagery: rotations yes, horizontal flip depends on orientation conventions.

Vision Transformers (ViT)

ViT applies the transformer architecture directly to images by splitting them into patches:

pythonCopy
# ViT conceptually: divide image into 16x16 patches, treat each as a "token"
# A 224x224 image → 14x14 = 196 patches

from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image

processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

image = Image.open("example.jpg")
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predicted_class = outputs.logits.argmax(-1).item()
print(model.config.id2label[predicted_class])

ViT advantages over CNNs: better global context (attention can connect distant parts of an image), better scaling with data. ViT disadvantages: requires more training data to work well (inductive biases of CNNs help with small datasets), more compute for the same accuracy on small datasets.

Current best practice: For small-to-medium datasets, use a pretrained CNN (EfficientNet, ResNet). For large datasets or when you need SOTA, use a pretrained ViT or hybrid (ConvNext).

Object Detection

Classification tells you what is in an image. Detection tells you where:

pythonCopy
import torch
import torchvision.transforms as T
from torchvision.models.detection import fasterrcnn_resnet50_fpn

# Load pretrained Faster R-CNN (COCO dataset: 80 classes)
model = fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

transform = T.Compose([T.ToTensor()])
image = transform(Image.open("street.jpg"))

with torch.no_grad():
    predictions = model([image])

# Filter by confidence threshold
threshold = 0.7
boxes = predictions[0]['boxes'][predictions[0]['scores'] > threshold]
labels = predictions[0]['labels'][predictions[0]['scores'] > threshold]
scores = predictions[0]['scores'][predictions[0]['scores'] > threshold]

print(f"Found {len(boxes)} objects with confidence > {threshold}")

For custom objects (your specific products, defects, etc.), fine-tune on your annotated dataset using a tool like Roboflow or Label Studio for annotation.

Production Considerations

Model size vs. latency:

Model	Parameters	ImageNet Accuracy	Typical Latency (CPU)
MobileNetV3-Small	2.5M	67.7%	~5ms
EfficientNet-B0	5.3M	77.1%	~10ms
ResNet-50	25M	80.9%	~30ms
ViT-B/16	86M	81.1%	~80ms

Export for production:

pythonCopy
# Export to ONNX for framework-agnostic serving
import torch.onnx

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["image"],
    output_names=["logits"],
    dynamic_axes={"image": {0: "batch_size"}}  # variable batch size
)

# Run with ONNX Runtime (2-4x faster than PyTorch on CPU)
import onnxruntime as ort
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
output = sess.run(None, {"image": input_array})

Evaluation Metrics

Task	Metric	Notes
Classification	Top-1 Accuracy, F1	Use F1 for class imbalance
Detection	[email protected], [email protected] .95	IoU threshold matters
Segmentation	mIoU	Mean intersection over union
Anomaly detection	AUC-ROC	Often no "normal" labeled examples

Common Mistakes

Not normalizing with ImageNet mean and std when using pretrained models. Pretrained CNNs (ResNet, EfficientNet, ViT) were trained on ImageNet-normalized inputs. If you feed raw [0, 255] or [0, 1] pixel values to a pretrained backbone, the activations at every layer will be out of distribution relative to what the weights expect, and fine-tuning will converge slowly or to a poor optimum. Always apply channel-wise normalization with the dataset statistics the model was originally trained on.

Augmenting the validation set. Data augmentation (random crop, flip, color jitter) is a regularization technique for training only. Applying it to validation images introduces randomness into your evaluation, meaning two runs with identical model weights can produce different validation accuracy numbers. Keep the validation pipeline deterministic: resize and normalize only.

Training from scratch on under-50K images instead of fine-tuning. Training a vision model from random initialization requires hundreds of thousands of labeled images to learn meaningful low-level features. Fine-tuning a pretrained backbone requires far less data because the feature extractor already understands edges, textures, and shapes. On small datasets, fine-tuning almost always outperforms training from scratch by a large margin.

What to Practice Next

• Fine-tune EfficientNet-B0 on a 10-class image dataset (e.g., a subset of CIFAR-100 or Food-101); run the same training with and without ImageNet normalization and compare final validation accuracy.
• Inspect five random training batches after augmentation and five validation batches after preprocessing; confirm the validation batches look like "clean" center-cropped images and not augmented ones.
• Visualize the learned filters of the first convolutional layer of a pretrained ResNet; compare them to filters learned when the same architecture is trained from scratch on 5K images.