| 49 | |
| 50 | |
| 51 | def track_memory_gpu(args, epoch): |
| 52 | log_dir = os.path.join(args.log_dir, args.exp_id) |
| 53 | os.makedirs(log_dir, exist_ok=True) |
| 54 | nvidia_smi.nvmlInit() |
| 55 | deviceCount = nvidia_smi.nvmlDeviceGetCount() |
| 56 | for i in range(deviceCount): |
| 57 | handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i) |
| 58 | info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle) |
| 59 | with open(os.path.join(log_dir, "TrackMemoryGPU.txt"), "a") as f: |
| 60 | f.write("\nEpoch: {} | Device {}: {}, Memory : ({:.2f}% free): {}(total), {} (free), {} (used)".format(epoch, i, nvidia_smi.nvmlDeviceGetName(handle), 100*info.free/info.total, info.total, info.free, info.used)) |
| 61 | nvidia_smi.nvmlShutdown() |