def trainNet(model, train_loader, val_loader, device, mask=None):
print("=" * 30)
# define the optimizer & learning rate
# optim = torch.optim.SGD(
# model.parameters(), lr=config["learning_rate"], weight_decay=0.0001, momentum=0.9, nesterov=True
# )
optim = torch.optim.AdamW(model.parameters(), lr=config["learning_rate"])
scheduler = StepLR(optim, step_size=config["lr_step_size"], gamma=config["lr_gamma"])
log_dir = "../runs/" + networkName + str(int(datetime.now().timestamp()))
if not os.path.exists(log_dir):
os.makedirs(log_dir)
writer = Visualizer(log_dir)
# Time for printing
training_start_time = time.time()
globaliter = 0
scheduler_count = 0
scaler = GradScaler()
# initialize the early_stopping object
early_stopping = EarlyStopping(log_dir, patience=config["patience"], verbose=True)
# Loop for n_epochs
for epoch in range(config["num_epochs"]):
writer.write_lr(optim, globaliter)
# train for one epoch
globaliter = train(model, train_loader, optim, device, writer, epoch, globaliter, scaler, mask)
# At the end of the epoch, do a pass on the validation set
val_loss = validate(model, val_loader, device, writer, epoch, mask)
# early_stopping needs the validation loss to check if it has decresed,
# and if it has, it will make a checkpoint of the current model
early_stopping(val_loss, model)
if early_stopping.early_stop:
print("Early stopping")
if scheduler_count == 2:
break
model.load_state_dict(torch.load(log_dir + "/checkpoint.pt"))
early_stopping.early_stop = False
early_stopping.counter = 0
scheduler.step()
scheduler_count += 1
if config["debug"] == True:
break
print("Training finished, took {:.2f}s".format(time.time() - training_start_time))
model.load_state_dict(torch.load(log_dir + "/checkpoint.pt"))
# remember to close the writer
writer.close()
def main():
# get settings from command line arguments
settings = CommandLineParser().parse_args()
# create problem
problem = GrayScottProblem(settings.size, coefficients=settings.coefficients)
# create visualizer
visualizer = Visualizer(problem.problem_size(), \
export=settings.export, \
keepalive=settings.keepalive, \
show=(not settings.noshow))
# create step generator
stop_point_generator = TimeStepper(settings.timesteps, \
settings.outputs, mode='linear')
# evolution loop
stop_point = next(stop_point_generator)
for step in range(settings.timesteps + 1):
# print progress message
if settings.verbose == True:
progress = 100 * step / settings.timesteps
print('{:3.0f}% finished'.format(progress), end='\r')
# trigger visualization
if step == stop_point:
visualizer.update(problem.v)
try:
stop_point = next(stop_point_generator)
except StopIteration:
pass
# evolve problem
problem.evolve()
else:
if settings.verbose == True:
print('\nEvolution finished')
visualizer.close()
class World:
def __init__(self, dt: float, width: float, height: float, ppm: float = 8):
self.dynamic_agents = []
self.static_agents = []
self.t = 0 # simulation time
self.dt = dt # simulation time step
self.visualizer = Visualizer(width, height, ppm=ppm)
def add(self, entity: Entity):
if entity.movable:
self.dynamic_agents.append(entity)
else:
self.static_agents.append(entity)
def tick(self):
for agent in self.dynamic_agents:
agent.tick(self.dt)
self.t += self.dt
def render(self):
self.visualizer.create_window(bg_color='gray')
self.visualizer.update_agents(self.agents)
@property
def agents(self):
return self.static_agents + self.dynamic_agents
def collision_exists(self, agent=None):
if agent is None:
for i in range(len(self.dynamic_agents)):
for j in range(i + 1, len(self.dynamic_agents)):
if self.dynamic_agents[
i].collidable and self.dynamic_agents[j].collidable:
if self.dynamic_agents[i].collidesWith(
self.dynamic_agents[j]):
return True
for j in range(len(self.static_agents)):
if self.dynamic_agents[
i].collidable and self.static_agents[j].collidable:
if self.dynamic_agents[i].collidesWith(
self.static_agents[j]):
return True
return False
if not agent.collidable: return False
for i in range(len(self.agents)):
if self.agents[i] is not agent and self.agents[
i].collidable and agent.collidesWith(self.agents[i]):
return True
return False
def close(self):
self.reset()
self.static_agents = []
if self.visualizer.window_created:
self.visualizer.close()
def reset(self):
self.dynamic_agents = []
self.t = 0