def create_train_recorder(self):
return TrainingRecorder(stuff_to_record=[
'total_reward', 'actor_loss', 'critic_loss', 'agent_colors',
'eval_episodes', 'eval_total_reward_explore',
'eval_total_reward_exploit'
])
# Create the optimizer and (potentially) load the optimizer state dictionary
optimizer = optim.SGD(network.parameters(), lr=args['learning_rate'], momentum=args['momentum'], weight_decay=args['weight_decay'])
if load_network_flag:
optimizer.load_state_dict(torch.load(os.path.abspath(f'{args["network_dir"]}{experiment}_optimizer.pth')))
# Create a learning rate scheduler -- this will reduce the learning rate by a factor when learning becomes stagnant
if args['scheduler'] == 'step':
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=35, gamma=0.1)
elif args['scheduler'] == 'bold':
scheduler = BoldDriver(optimizer, network, decrease_factor=0.5, increase_factor=1.2)
# Create classification loss function
classification_loss_func = nn.CrossEntropyLoss()
# Create a recorder
recorder = TrainingRecorder(args['results_dir'], args['network_dir'], parameters, new=(parameters.epoch == 0))
# Begin training the neural network
for epoch in range(parameters.epoch, parameters.num_epochs):
true_classes = np.empty(0)
predicted_classes = np.empty(0)
running_loss = 0.0
network.train()
for batch_num, batch_sample in enumerate(train_dataloader):
optimizer.zero_grad()
# Load in batch image data
image_data = batch_sample[0]
image_data.requires_grad = False