File size: 3,623 Bytes
b678162 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import torch
import sys
from pathlib import Path
# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent))
def check_model_memory():
# Simple config object
class SimpleConfig:
class ModelConfig:
decoder_depth = 4
model = ModelConfig()
cfg = SimpleConfig()
# Import after path is set
from dpm.model import VDPM
# Create model on CPU first to count parameters
print("Creating model...")
model = VDPM(cfg)
# Count parameters
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"\n{'='*60}")
print(f"MODEL SIZE ANALYSIS FOR RTX 3070 Ti (8GB)")
print(f"{'='*60}")
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")
print(f"\nEstimated model weights memory:")
print(f" - FP32 (float32): {total_params * 4 / 1024**3:.2f} GB")
print(f" - FP16 (float16): {total_params * 2 / 1024**3:.2f} GB")
print(f" - BF16 (bfloat16): {total_params * 2 / 1024**3:.2f} GB")
print(f" - INT8 (quantized): {total_params * 1 / 1024**3:.2f} GB <-- RECOMMENDED for 8GB GPU")
# Estimate activation memory for typical input
batch_size = 1
num_frames = 5 # typical video length
img_size = 518
print(f"\nEstimated activation memory (batch={batch_size}, frames={num_frames}, img_size={img_size}):")
# Input images: [B, S, 3, H, W]
input_mem = batch_size * num_frames * 3 * img_size * img_size * 4 / 1024**3
print(f" - Input images (FP32): {input_mem:.2f} GB")
# Rough estimate for activations (can be 2-4x model size during forward pass)
activation_mem_estimate = total_params * 2 * 3 / 1024**3 # conservative estimate
print(f" - Activations (estimate): {activation_mem_estimate:.2f} GB")
# Calculate total for different precision modes
total_fp16 = (total_params * 2 / 1024**3) + input_mem + activation_mem_estimate
total_int8 = (total_params * 1 / 1024**3) + input_mem + (activation_mem_estimate * 0.6) # INT8 reduces activations too
print(f"\nTotal estimated GPU memory needed:")
print(f" - With FP16/BF16: {total_fp16:.2f} GB")
print(f" - With INT8 quantization: {total_int8:.2f} GB <-- FITS IN 8GB!")
print(f"Your RTX 3070 Ti has: 8 GB VRAM")
if total_int8 <= 8:
print(f"\n✓ With INT8 quantization, model will fit in GPU memory!")
print(f" Set USE_QUANTIZATION = True in gradio_demo.py")
elif total_fp16 > 8:
print(f"\n⚠️ WARNING: Even with INT8 ({total_int8:.2f} GB), memory is tight")
print(f" Recommendations:")
print(f" 1. Use INT8 quantization (USE_QUANTIZATION = True)")
print(f" 2. Reduce number of input frames to {num_frames} or fewer")
print(f" 3. Clear CUDA cache between batches")
else:
print(f"\n✓ Model should fit with FP16!")
print(f"{'='*60}\n")
# Check actual GPU memory if CUDA available
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Total GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
print(f"Current GPU memory allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"Current GPU memory cached: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
if __name__ == "__main__":
check_model_memory()
|