Module neuralops.trainer.trainer_svd

Functions

def benchmark_model(model: torch.nn.modules.module.Module,
data_loader: torch.utils.data.dataloader.DataLoader,
device: torch.device) ‑> dict

Expand source code

def benchmark_model(model: nn.Module, data_loader: DataLoader, device: torch.device) -> dict:
    """Benchmark model."""
    model.eval()
    memory_before = psutil.Process().memory_info().rss / 1024 / 1024
    start_time = time.time()
    batch_times = []
    with torch.no_grad():
        for batch in data_loader:
            b_start = time.time()
            if batch["input_ids"].dim() == 3:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                for i in range(input_ids.shape[1]):
                    _ = model(
                        input_ids=input_ids[:, i, :],
                        attention_mask=attention_mask[:, i, :],
                    )
            else:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                _ = model(input_ids=input_ids, attention_mask=attention_mask)
            if device.type == "cuda":
                torch.cuda.synchronize()
            batch_times.append(time.time() - b_start)
    total_time = time.time() - start_time
    memory_after = psutil.Process().memory_info().rss / 1024 / 1024
    return {
        "total_time": total_time,
        "avg_batch_time": np.mean(batch_times) if batch_times else 0,
        "samples_per_second": (len(data_loader.dataset) / total_time if total_time > 0 else 0),
        "memory_usage_mb": memory_after - memory_before,
        "peak_memory_mb": memory_after,
        "throughput": len(data_loader.dataset) / total_time if total_time > 0 else 0,
    }

Benchmark model.

def count_parameters(model: torch.nn.modules.module.Module, trainable_only: bool = True) ‑> int

Expand source code

def count_parameters(model: torch.nn.Module, trainable_only: bool = True) -> int:
    """Count parameters."""
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())

Count parameters.

def train_model(model: torch.nn.modules.module.Module,
train_loader: torch.utils.data.dataloader.DataLoader,
val_loader: torch.utils.data.dataloader.DataLoader,
device: torch.device,
config,
num_labels: int,
is_multichoice: bool = False) ‑> Any

Expand source code

def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    device: torch.device,
    config,
    num_labels: int,
    is_multichoice: bool = False,
) -> Any:
    """Train model."""
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
    steps_per_epoch = len(train_loader)
    for epoch in range(config.epochs):
        epoch_loss = 0.0
        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{config.epochs}", leave=False)
        for batch_idx, batch in enumerate(pbar):
            optimizer.zero_grad(set_to_none=True)
            if is_multichoice:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                true_labels = batch["labels"].to(device)
                batch_size, num_choices, seq_len = input_ids.shape
                input_ids_flat = input_ids.view(-1, seq_len)
                attention_mask_flat = attention_mask.view(-1, seq_len)
                outputs = model(input_ids=input_ids_flat, attention_mask=attention_mask_flat)
                logits = outputs["logits"]
                combined_labels = true_labels.repeat_interleave(num_choices)
                loss = nn.functional.cross_entropy(logits, combined_labels)
            else:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)
                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                logits = outputs["logits"]
                loss = nn.functional.cross_entropy(logits, labels)
            if torch.isfinite(loss):
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()
                epoch_loss += loss.item()
            pbar.set_postfix({"loss": f"{loss.item():.4f}"})
    val_accuracy = evaluate_model(
        model, val_loader, device, num_labels, is_multichoice=is_multichoice
    )
    return {"final_accuracy": val_accuracy, "avg_loss": epoch_loss / steps_per_epoch}

Train model.