153 lines
4.4 KiB
Python
Executable File
153 lines
4.4 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""
|
|
Quick model testing script to verify model creation and inference.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import torch
|
|
|
|
# Add project root to the path to enable imports
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from models.detection import get_maskrcnn_model
|
|
from utils.data_utils import PennFudanDataset, get_transform
|
|
|
|
|
|
def test_model_creation():
|
|
"""Test that we can create the model."""
|
|
print("Testing model creation...")
|
|
model = get_maskrcnn_model(
|
|
num_classes=2, pretrained=False, pretrained_backbone=False
|
|
)
|
|
print("✓ Model created successfully")
|
|
return model
|
|
|
|
|
|
def test_model_forward(model, device):
|
|
"""Test model forward pass with random inputs."""
|
|
print("\nTesting model forward pass...")
|
|
|
|
# Create a random batch
|
|
image = torch.rand(3, 300, 400, device=device) # Random image
|
|
|
|
# Create a random target
|
|
target = {
|
|
"boxes": torch.tensor(
|
|
[[100, 100, 200, 200]], dtype=torch.float32, device=device
|
|
),
|
|
"labels": torch.tensor([1], dtype=torch.int64, device=device),
|
|
"masks": torch.randint(0, 2, (1, 300, 400), dtype=torch.uint8, device=device),
|
|
"image_id": torch.tensor([0], device=device),
|
|
"area": torch.tensor([10000.0], dtype=torch.float32, device=device),
|
|
"iscrowd": torch.tensor([0], dtype=torch.uint8, device=device),
|
|
}
|
|
|
|
# Test inference mode (no targets)
|
|
model.eval()
|
|
with torch.no_grad():
|
|
start_time = time.time()
|
|
output_inference = model([image])
|
|
inference_time = time.time() - start_time
|
|
|
|
# Verify inference output
|
|
print(f"✓ Inference mode output: {type(output_inference)}")
|
|
print(f"✓ Inference time: {inference_time:.3f}s")
|
|
print(f"✓ Detection boxes shape: {output_inference[0]['boxes'].shape}")
|
|
print(f"✓ Detection scores shape: {output_inference[0]['scores'].shape}")
|
|
|
|
# Test training mode (with targets)
|
|
model.train()
|
|
start_time = time.time()
|
|
output_train = model([image], [target])
|
|
train_time = time.time() - start_time
|
|
|
|
# Verify training output
|
|
print(f"✓ Training mode output: {type(output_train)}")
|
|
print(f"✓ Training time: {train_time:.3f}s")
|
|
|
|
# Print loss values
|
|
for loss_name, loss_value in output_train.items():
|
|
print(f"✓ {loss_name}: {loss_value.item():.4f}")
|
|
|
|
return output_train
|
|
|
|
|
|
def test_model_backward(model, loss_dict, device):
|
|
"""Test model backward pass."""
|
|
print("\nTesting model backward pass...")
|
|
|
|
# Calculate total loss
|
|
total_loss = sum(loss for loss in loss_dict.values())
|
|
|
|
# Create optimizer
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
|
|
|
|
# Backward pass
|
|
start_time = time.time()
|
|
optimizer.zero_grad()
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
backward_time = time.time() - start_time
|
|
|
|
print("✓ Backward pass and optimization completed")
|
|
print(f"✓ Backward time: {backward_time:.3f}s")
|
|
|
|
# Check that gradients were calculated
|
|
has_gradients = any(
|
|
param.grad is not None for param in model.parameters() if param.requires_grad
|
|
)
|
|
print(f"✓ Model has gradients: {has_gradients}")
|
|
|
|
|
|
def test_dataset():
|
|
"""Test that we can load the dataset."""
|
|
print("\nTesting dataset loading...")
|
|
|
|
data_root = "data/PennFudanPed"
|
|
if not os.path.exists(data_root):
|
|
print("✗ Dataset not found at", data_root)
|
|
return None
|
|
|
|
# Create dataset
|
|
dataset = PennFudanDataset(root=data_root, transforms=get_transform(train=True))
|
|
print(f"✓ Dataset loaded with {len(dataset)} samples")
|
|
|
|
# Test loading a sample
|
|
start_time = time.time()
|
|
img, target = dataset[0]
|
|
load_time = time.time() - start_time
|
|
|
|
print(f"✓ Sample loaded in {load_time:.3f}s")
|
|
print(f"✓ Image shape: {img.shape}")
|
|
print(f"✓ Target boxes shape: {target['boxes'].shape}")
|
|
|
|
return dataset
|
|
|
|
|
|
def main():
|
|
"""Run all tests."""
|
|
print("=== Quick Model Testing Script ===")
|
|
|
|
# Set device
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {device}")
|
|
|
|
# Run tests
|
|
model = test_model_creation()
|
|
model.to(device)
|
|
|
|
loss_dict = test_model_forward(model, device)
|
|
|
|
test_model_backward(model, loss_dict, device)
|
|
|
|
test_dataset()
|
|
|
|
print("\n=== All tests completed successfully ===")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|