Spaces:
Sleeping
Sleeping
| import os | |
| from datetime import datetime | |
| import socket | |
| import torch | |
| import sys | |
| from torch.utils.tensorboard import SummaryWriter | |
| import csv | |
| class Logger(object): | |
| def __init__(self, directory, write, save_freq=10, comment=""): | |
| self.parent_dir = directory | |
| self.write = write | |
| self.dir = "" | |
| self.fig_index = 0 | |
| self.model_index = 0 | |
| self.save_freq = save_freq | |
| if self.write: | |
| self.boardWriter = SummaryWriter(comment=comment) | |
| self.dir = self.boardWriter.log_dir | |
| self.log(f"Logs from {self.dir}\n{' '.join(sys.argv)}\n") | |
| def write_to_board(self, name, scalars, index=0): | |
| self.log(f"{name} at {index}: {str(scalars)}") | |
| if self.write: | |
| for key, value in scalars.items(): | |
| self.boardWriter.add_scalar(f"{name}/{key}", value, index) | |
| def plot_res(self, losses, distances): | |
| if len(losses) == 0 or not self.write: | |
| return | |
| import matplotlib.pyplot as plt | |
| fig, axs = plt.subplots(2) | |
| axs[0].plot(list(range(len(losses))), losses, color='orange') | |
| axs[0].set_xlabel("Steps") | |
| axs[0].set_ylabel("Loss") | |
| axs[0].set_title("Training") | |
| axs[0].set_yscale('log') | |
| for dist in distances: | |
| axs[1].plot(list(range(len(dist))), dist) | |
| axs[1].set_xlabel("Steps") | |
| axs[1].set_ylabel("Distance change") | |
| axs[1].set_title("Training") | |
| if self.fig_index > 0: | |
| os.remove(os.path.join(self.dir, f"res{self.fig_index-1}.png")) | |
| fig.savefig(os.path.join(self.dir, f"res{self.fig_index}.png")) | |
| self.boardWriter.add_figure(f"res{self.fig_index}", fig) | |
| self.fig_index += 1 | |
| def log(self, message, step=0): | |
| print(str(message)) | |
| if self.write: | |
| # self.boardWriter.add_text("log", str(message), step) | |
| with open(os.path.join(self.dir, "logs.txt"), "a") as logs: | |
| logs.write(str(message) + "\n") | |
| def save_model(self, state_dict, name="dqn.pt", forced=False): | |
| if not self.write: | |
| return | |
| if (forced or | |
| (self.model_index > 0 and self.model_index % self.save_freq == 0)): | |
| torch.save(state_dict, os.path.join(self.dir, name)) | |
| def write_locations(self, row): | |
| self.log(str(row)) | |
| if self.write: | |
| with open(os.path.join(self.dir, 'results.csv'), | |
| mode='a', newline='') as f: | |
| res_writer = csv.writer(f) | |
| res_writer.writerow(row) | |