def make(config):
model = nn.DataParallel(STLSTM(1, config['hidden_size'])).to(device)
print(config['depths'])
if config['dataset'] == 'regimeA':
root = config['data_folder_regimeA']
elif config['dataset'] == 'regimeB':
root = config['data_folder_regimeB']
#train_dataset = BarkleyDataset(root=root,
# train=True,
# depths=config['depths'],
# time_steps=config['time_steps'])
#n_train = int(len(train_dataset)*0.90+0.5)
#n_val = int(len(train_dataset)*0.10+0.5)
#torch.manual_seed(42)
#train_dataset, test_dataset = random_split(train_dataset, [n_train, n_val])
#test_dataset.train = False
#train_loader = DataLoader(train_dataset, config['batch_size'], shuffle=True, num_workers=0, pin_memory=True, drop_last=True)
#test_loader = DataLoader(test_dataset, 2, shuffle=True, num_workers=0, pin_memory=True, drop_last=True)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
return model, criterion, optimizer
def getModel(cfg, suffix):
model = nn.DataParallel(STLSTM(1,64))
prefix = 'STLSTM_'
model_file = prefix + suffix
folder = f'../models/{cfg.dataset}'
#model_file = f'STLSTM_t{cfg.time_step}_d{cfg.depth}'
file = os.path.join(folder, model_file)
print(file)
model.load_state_dict(torch.load(file, map_location=device), strict=True)
return model
def getModel(cfg, **kwargs):
model = nn.DataParallel(STLSTM(1,64))
suffix = kwargs.pop('suffix', "[0]")
folder = f'../models/regime{cfg.regime}'
model_file = f'{cfg.prefix}{suffix}'
#model_file = f'STLSTM_t{cfg.time_step}_d{cfg.depth}'
file = os.path.join(folder, model_file)
model.load_state_dict(torch.load(file, map_location=device), strict=True)
return model
def getModels(cfg, suffixes):
models = []
for suffix in suffixes:
model = nn.DataParallel(STLSTM(1,64))
folder = f'../models/{cfg.dataset}'
folder = '/home/roland/Projekte/FromSurface2Depth/models/regimeB'
model_file = f'STLSTM_{suffix}'
file = os.path.join(folder, model_file)
print(file)
model.load_state_dict(torch.load(file, map_location=device), strict=True)
models.append(model)
return models
def make(config):
model = nn.DataParallel(STLSTM(1,config['num_features'])).to(device)
if config['load_weights']:
try:
model.load_state_dict(torch.load(config['weights_file'], map_location=device), strict=True)
except FileNotFoundError:
print('File not found, initialize new weights')
train_dataset = BarkleyDataset(root=config['dataset_dir'],
train=True,
depth=config['depth'],
time_steps=config['time_step'])
n_train = int(len(train_dataset)*0.90+0.5)
n_val = int(len(train_dataset)*0.10+0.5)
torch.manual_seed(42)
train_dataset, test_dataset = random_split(train_dataset, [n_train, n_val])
test_dataset.train = False
train_loader = DataLoader(train_dataset, config['batch_size'], shuffle=True, num_workers=0, pin_memory=True, drop_last=True)
test_loader = DataLoader(test_dataset, 2, shuffle=True, num_workers=0, pin_memory=True, drop_last=True)
if config['loss_fn'] =='MSE':
criterion = nn.MSELoss()
elif config['loss_fn']=='MAE':
criterion = nn.L1Loss()
if config['optimizer'] =='Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
elif config['optimizer']=='SGD':
optimizer = torch.optim.SGD(model.parameters(), lr=config['lr'])
return model, train_loader, test_loader, criterion, optimizer