def save_test_results(mat_prop, classification_list):
# Load up a saved network.
model = Model(CrabNet(compute_device=compute_device).to(compute_device))
model.load_network(f'{mat_prop}.pth')
if mat_prop in classification_list:
model.classification = True
# Load the data you want to predict with
test_data = rf'data\benchmark_data\{mat_prop}\test.csv'
model.load_data(test_data) # data is reloaded to model.data_loader
output = model.predict(model.data_loader) # predict the data saved here
if model.classification:
auc = roc_auc_score(output[0], output[1])
print(f'\n{mat_prop} ROC AUC: {auc:0.3f}')
else:
print(f'\n{mat_prop} mae: {abs(output[0] - output[1]).mean():0.3f}')
# save your predictions to a csv
save_results(output, f'{mat_prop}_output.csv')
from utils.utils import CONSTANTS
compute_device = get_compute_device()
# %%
mat_prop = 'mp_bulk_modulus'
crabnet_params = {'d_model': 512, 'N': 3, 'heads': 4}
model = Model(CrabNet(**crabnet_params, compute_device=compute_device).to(compute_device))
model.load_network(f'{mat_prop}.pth')
# Load the data you want to predict with
test_data = rf'data\benchmark_data\{mat_prop}\train.csv'
model.load_data(test_data) # data is reloaded to model.data_loader
output = model.predict(model.data_loader) # predict the data saved here
# %%
class SaveOutput:
def __init__(self):
self.outputs = []
def __call__(self, module, module_in, module_out):
self.outputs.append(module_out)
def clear(self):
self.outputs = []
save_output = SaveOutput()
hook_handles = []
CrabNet(**torchnet_params,
compute_device=compute_device).to(compute_device))
model.load_network(f'{mat_prop}.pth')
hook_handles = []
# Insert forward hooks into model
for layer in model.model.modules():
if isinstance(layer, torch.nn.modules.activation.MultiheadAttention):
# print('isinstance')
handle = layer.register_forward_hook(save_output)
hook_handles.append(handle)
model.load_data(data) # data is reloaded to model.data_loader
save_output.clear()
output = model.predict(model.data_loader)
elem_pred = [(i, out, output[1][i]) for i, out in enumerate(output[2])
if elem_sym in out]
df_elem = pd.DataFrame(elem_pred, columns=['idx', 'formula', 'prediction'])
mod_out = save_output.outputs
n_mats = len(mod_out) # number of output matrices from hook
bsz = model.data_loader.batch_size # batch size from data loader
B = len(model.data_loader) # total number of batches from data loader
H = model.model.heads # number of heads
N = model.model.N # number of layers
n_data = len(model.data_loader.dataset)
n_elements = model.n_elements