| 112 | ax.text(x0, y0, label, color=color_mask, fontsize='large', fontfamily='sans-serif') |
| 113 | |
| 114 | def get_gpu_memory(max_gpus=None): |
| 115 | gpu_memory = [] |
| 116 | num_gpus = ( |
| 117 | torch.cuda.device_count() |
| 118 | if max_gpus is None |
| 119 | else min(max_gpus, torch.cuda.device_count()) |
| 120 | ) |
| 121 | |
| 122 | for gpu_id in range(num_gpus): |
| 123 | with torch.cuda.device(gpu_id): |
| 124 | device = torch.cuda.current_device() |
| 125 | gpu_properties = torch.cuda.get_device_properties(device) |
| 126 | total_memory = gpu_properties.total_memory / (1024 ** 3) |
| 127 | allocated_memory = torch.cuda.memory_allocated() / (1024 ** 3) |
| 128 | available_memory = total_memory - allocated_memory |
| 129 | gpu_memory.append(available_memory) |
| 130 | return gpu_memory |
| 131 | |
| 132 | def load_model( |
| 133 | model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, lora_weights=None |