| 3 | import nvidia_smi |
| 4 | |
| 5 | def save_logs( |
| 6 | args, |
| 7 | train_losses, |
| 8 | train_ppls, |
| 9 | train_times, |
| 10 | valid_losses, |
| 11 | valid_ppls, |
| 12 | valid_times, |
| 13 | test_loss, |
| 14 | test_ppl, |
| 15 | test_time, |
| 16 | tot_params, |
| 17 | tot_learnable_params, |
| 18 | ): |
| 19 | log_dir = os.path.join(args.log_dir, args.exp_id) |
| 20 | os.makedirs(log_dir, exist_ok=True) |
| 21 | |
| 22 | # Log arguments |
| 23 | with open(os.path.join(log_dir, "args.json"), "w") as f: |
| 24 | json.dump(args.__dict__, f, indent=2) |
| 25 | |
| 26 | # Log train, val, and test losses and perplexities |
| 27 | with open(os.path.join(log_dir, "train_loss.txt"), "w") as f: |
| 28 | f.write("\n".join(str(item) for item in train_losses)) |
| 29 | with open(os.path.join(log_dir, "train_ppl.txt"), "w") as f: |
| 30 | f.write("\n".join(str(item) for item in train_ppls)) |
| 31 | with open(os.path.join(log_dir, "train_time.txt"), "w") as f: |
| 32 | f.write("\n".join(str(item) for item in train_times)) |
| 33 | with open(os.path.join(log_dir, "valid_loss.txt"), "w") as f: |
| 34 | f.write("\n".join(str(item) for item in valid_losses)) |
| 35 | with open(os.path.join(log_dir, "valid_ppl.txt"), "w") as f: |
| 36 | f.write("\n".join(str(item) for item in valid_ppls)) |
| 37 | with open(os.path.join(log_dir, "valid_time.txt"), "w") as f: |
| 38 | f.write("\n".join(str(item) for item in valid_times)) |
| 39 | with open(os.path.join(log_dir, "test_loss.txt"), "w") as f: |
| 40 | f.write(f"{test_loss}\n") |
| 41 | with open(os.path.join(log_dir, "test_ppl.txt"), "w") as f: |
| 42 | f.write(f"{test_ppl}\n") |
| 43 | with open(os.path.join(log_dir, "test_time.txt"), "w") as f: |
| 44 | f.write(f"{test_time}\n") |
| 45 | with open(os.path.join(log_dir, "tot_params.txt"), "w") as f: |
| 46 | f.write(f"{tot_params}\n") |
| 47 | with open(os.path.join(log_dir, "tot_learnable_params.txt"), "w") as f: |
| 48 | f.write(f"{tot_learnable_params}\n") |
| 49 | |
| 50 | |
| 51 | def track_memory_gpu(args, epoch): |