Пример #1
0
 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'
     ])
Пример #2
0
    # 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