def test_net():
# Function made only to test a pretrained network.
config = json.load(open('./config.json'))
config['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
model = EarthNetComplex().to(config['device'])
checkpoint = torch.load(os.path.join(config['exp_path'],
config['exp_name'],
'latest_checkpoint.pkl'),
map_location=config['device'])
model.load_state_dict(checkpoint['model_weights'])
criterion = getattr(loss_functions, config['loss_function'])
data_manager = DataManager(config)
_, _, test_loader = data_manager.get_train_eval_test_dataloaders()
trainer = Trainer(model, None, None, criterion, None, None, config)
trainer.test_net(test_loader)
def main():
config = json.load(open('./config.json'))
config['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
os.mkdir(os.path.join(config['exp_path'], config['exp_name']))
except FileExistsError:
print("Director already exists! It will be overwritten!")
model = EarthNetComplex().to(config['device'])
model.apply(EarthNetComplex.init_weights)
# Save info about experiment
save_logs_about(os.path.join(config['exp_path'], config['exp_name']),
json.dumps(config, indent=2))
shutil.copy(model.get_path(),
os.path.join(config['exp_path'], config['exp_name']))
criterion = getattr(loss_functions, config['loss_function'])
optimizer = optim.Adam(model.parameters(),
lr=config['lr'],
weight_decay=config['weight_decay'])
lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
config['lr_sch_step'],
gamma=config['lr_sch_gamma'],
last_epoch=-1)
data_manager = DataManager(config)
train_loader, validation_loader, test_loader = data_manager.get_train_eval_test_dataloaders(
)
trainer = Trainer(model, train_loader, validation_loader, criterion,
optimizer, lr_scheduler, config)
trainer.train()
trainer.test_net(test_loader)