def train_npi(
steplists,
epochs: int = 100,
batch_size: int = 1,
pretrained_encoder_weights: str = f'{config.outdir}/weights/f_enc.weights'
):
state_encoder = StateEncoder().to(config.device)
warm_up: list = list(
filter(
lambda question: 0 <= question['question'][0] < 100 and 0 <=
question['question'][1] < 100, steplists))
if not os.path.exists(pretrained_encoder_weights):
print('start trainning f_enc model')
train_f_enc(warm_up, epochs=epochs, batch_size=batch_size)
else:
state_encoder.load_state_dict(torch.load(pretrained_encoder_weights))
state_encoder.trainable = False
npi = NPI(state_encoder).to(config.device)
optimizer = optim.Adam(npi.parameters(), lr=1e-4, weight_decay=1e-6)
# warm up with single digit addjj
for _ in range(10):
if train_with_plot(npi, optimizer, warm_up, epochs=100):
break
while True:
if train_with_plot(npi,
optimizer,
steplists,
epochs=100,
skip_correct=False):
break
train_data = pickle.load(f)
EVAL_DATA_PATH = 'tasks/reverse_polish/data/eval_8.pik'
with open(EVAL_DATA_PATH, 'rb', ) as f:
eval_data = pickle.load(f)
TEST_DATA_PATH = 'tasks/reverse_polish/data/test_8.pik'
with open(TEST_DATA_PATH, 'rb', ) as f:
test_data = pickle.load(f)
func_core = RevPolishCore().to(device)
npi = NPI(func_core, CONFIG).to(device)
print_net(npi)
optimizer = torch.optim.Adam(npi.parameters(), lr=1e-4)
# optimizer = torch.optim.SGD(npi.parameters(), lr=1e-4)
lr_schedulers = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_num_epochs)
if not FLG_TEST:
for epoch in range(start_epoch, max_num_epochs+1):
run_epoch(npi, 'train', train_data, writer, os.path.join(exp_dir, 'npi_{}.pth'.format(epoch)))
lr_schedulers.step(epoch)
if epoch % 2 == 0:
def_loss, total_loss, ter_accs, pro_accs = run_epoch(npi, 'eval', eval_data, writer, os.path.join(exp_dir, 'npi_{}.pth'.format(epoch)))
if def_loss < Best_results['def_loss']:
Best_results['def_loss'] = def_loss
EVAL_DATA_PATH = 'tasks/reverse_polish/data/eval.pik'
with open(EVAL_DATA_PATH, 'rb', ) as f:
eval_data = pickle.load(f)
TEST_DATA_PATH = 'tasks/reverse_polish/data/test.pik'
with open(TEST_DATA_PATH, 'rb', ) as f:
test_data = pickle.load(f)
func_core = RevPolishCore().to(device)
npi = NPI(func_core, CONFIG).to(device)
print_net(npi)
# optimizer = torch.optim.Adam(npi.parameters(), lr=1e-4)
optimizer = torch.optim.SGD(npi.parameters(), lr=1e-3)
lr_schedulers = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_num_epochs)
if not FLG_TEST:
for epoch in range(start_epoch, max_num_epochs+1):
run_epoch(npi, 'train', train_data, writer, os.path.join(exp_dir, 'npi_{}.pth'.format(epoch)))
lr_schedulers.step(epoch)
if epoch % 2 == 0:
def_loss, total_loss, ter_accs, pro_accs = run_epoch(npi, 'eval', eval_data, writer, os.path.join(exp_dir, 'npi_{}.pth'.format(epoch)))
if def_loss < Best_results['def_loss']:
Best_results['def_loss'] = def_loss
Best_results['epoch_def_loss'] = epoch
