cudnn.benchmark = True #May train faster but cost more memory
# define loss function
criterion = nn.MSELoss()
# get the data
train_loader = data.create_dataset_maps('train', seed, f_in, f_out, batch_size,
verbose=True)
valid_loader = data.create_dataset_maps('valid', seed, f_in, f_out, batch_size,
verbose=True)
test_loader = data.create_dataset_maps('test', seed, f_in, f_out, batch_size,
verbose=True)
# define the model and get total number of parameters
if model=='model_a': model = architecture.model_a(hidden).to(device)
elif model=='model_b': model = architecture.model_b(hidden).to(device)
elif model=='model_c': model = architecture.model_c(hidden).to(device)
elif model=='model_d': model = architecture.model_d(hidden).to(device)
elif model=='model_e': model = architecture.model_e(hidden).to(device)
elif model=='model_f': model = architecture.model_f(hidden).to(device)
else: raise Exception('model not supported')
network_total_params = sum(p.numel() for p in model.parameters())
print('total number of parameters in the model: %d'%network_total_params)
# define the optimizer
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=wd)
# get validation loss
if os.path.exists(f_model): model.load_state_dict(torch.load(f_model))
print('Computing initial validation loss')
model.eval()
cudnn.benchmark = True #May train faster but cost more memory
# define loss function
criterion = nn.MSELoss()
# get the data
train_loader = data.create_dataset_cubes('train', seed, f_in, f_out, batch_size,
verbose=True)
valid_loader = data.create_dataset_cubes('valid', seed, f_in, f_out, batch_size,
verbose=True)
test_loader = data.create_dataset_cubes('test', seed, f_in, f_out, batch_size,
verbose=True)
# define the model
#model = architecture.model_a(hidden).to(device)
model = architecture.model_b(hidden).to(device)
# load best-models, if they exists
if os.path.exists(f_model): model.load_state_dict(torch.load(f_model,
map_location=torch.device(device)))
# define the optimizer
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=wd)
# get the initial validation loss
print('Computing initial validation loss')
model.eval()
min_valid_loss, num_points = 0.0, 0
for x,y in valid_loader:
with torch.no_grad():