Spaces:
Sleeping
Sleeping
| import collections | |
| import sys | |
| import numpy as np | |
| import pandas as pd | |
| import random | |
| import torch | |
| import time | |
| import os | |
| import json | |
| import tifffile | |
| import h3 | |
| import setup | |
| from sklearn.linear_model import RidgeCV | |
| from sklearn.preprocessing import MinMaxScaler | |
| from torch.utils.data import Subset | |
| import utils | |
| import models | |
| import datasets | |
| from calendar import monthrange | |
| from torch.nn.functional import logsigmoid, softmax | |
| import torch.nn as nn | |
| from tqdm import tqdm | |
| import csv | |
| def format_tensor(tensor): | |
| # Convert tensor to list, then flatten to string | |
| tensor_list = tensor.tolist() # Converts the tensor to a Python list | |
| return str(tensor_list).replace('\n', '').replace(' ', '') | |
| class EvaluatorSNT: | |
| def __init__(self, train_params, eval_params): | |
| self.train_params = train_params | |
| self.eval_params = eval_params | |
| with open('paths.json', 'r') as f: | |
| paths = json.load(f) | |
| D = np.load(os.path.join(paths['snt'], 'snt_res_5.npy'), allow_pickle=True) | |
| D = D.item() | |
| self.loc_indices_per_species = D['loc_indices_per_species'] | |
| self.labels_per_species = D['labels_per_species'] | |
| self.taxa = D['taxa'] | |
| self.obs_locs = D['obs_locs'] | |
| self.obs_locs_idx = D['obs_locs_idx'] | |
| self.pos_eval_data_loc = os.path.join(paths['data'], 'positive_eval_data.npz') | |
| self.background_eval_data_loc = os.path.join(paths['data'], '10000_background_negs.npz') | |
| def get_labels(self, species): | |
| species = str(species) | |
| lat = [] | |
| lon = [] | |
| gt = [] | |
| for hx in self.data: | |
| cur_lat, cur_lon = h3.h3_to_geo(hx) | |
| if species in self.data[hx]: | |
| cur_label = int(len(self.data[hx][species]) > 0) | |
| gt.append(cur_label) | |
| lat.append(cur_lat) | |
| lon.append(cur_lon) | |
| lat = np.array(lat).astype(np.float32) | |
| lon = np.array(lon).astype(np.float32) | |
| obs_locs = np.vstack((lon, lat)).T | |
| gt = np.array(gt).astype(np.float32) | |
| return obs_locs, gt | |
| def run_evaluation(self, model, enc, extra_input=None): | |
| results = {} | |
| # set seeds: | |
| np.random.seed(self.eval_params['seed']) | |
| random.seed(self.eval_params['seed']) | |
| # evaluate the geo model for each taxon | |
| results['per_species_average_precision_all'] = np.zeros((len(self.taxa)), dtype=np.float32) | |
| # get eval locations and apply input encoding | |
| obs_locs = torch.from_numpy(self.obs_locs).to(self.eval_params['device']) | |
| loc_feat = torch.cat([enc.encode(obs_locs), extra_input.expand(obs_locs.shape[0], -1)], dim=1) if extra_input is not None else enc.encode(obs_locs) | |
| # get classes to eval | |
| classes_of_interest = torch.zeros(len(self.taxa), dtype=torch.int64) | |
| for tt_id, tt in enumerate(self.taxa): | |
| class_of_interest = np.where(np.array(self.train_params['class_to_taxa']) == tt)[0] | |
| if len(class_of_interest) != 0: | |
| classes_of_interest[tt_id] = torch.from_numpy(class_of_interest) | |
| if self.eval_params['extract_pos']: | |
| assert 'HyperNet' in self.train_params['model'] | |
| model = model.pos_enc | |
| self.train_params['model'] = 'ResidualFCNet' | |
| if ('CombinedModel' in self.train_params['model']) or ('MultiInputModel' in self.train_params['model']): | |
| with torch.no_grad(): | |
| dummy_context_mask = None | |
| dummy_context_sequence = None | |
| # generate model predictions for classes of interest at eval locations | |
| loc_emb = model(x=loc_feat, context_sequence=dummy_context_sequence, context_mask=dummy_context_mask, | |
| class_ids=classes_of_interest, return_feats=True) | |
| classes_of_interest = classes_of_interest.to(self.eval_params["device"]) | |
| wt = model.get_eval_embeddings(classes_of_interest) | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif self.train_params['model'] == 'VariableInputModel': | |
| with torch.no_grad(): | |
| loc_emb = model.get_loc_emb(x=loc_feat) | |
| classes_of_interest = classes_of_interest.to(self.eval_params["device"]) | |
| wt = model.get_eval_embeddings(classes_of_interest) | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif 'HyperNet' not in self.train_params['model'] and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| with torch.no_grad(): | |
| # generate model predictions for classes of interest at eval locations | |
| loc_emb = model(loc_feat, return_feats=True) | |
| wt = model.class_emb.weight[classes_of_interest, :] | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| if self.train_params['model'] == 'ResidualFCNet': | |
| import datasets | |
| # from sklearn.linear_model import LogisticRegression | |
| # with open('paths.json', 'r') as f: | |
| # paths = json.load(f) | |
| # data_dir = paths['train'] | |
| # obs_file = os.path.join(data_dir, self.train_params['obs_file']) | |
| # taxa_file = os.path.join(data_dir, self.train_params['taxa_file']) | |
| # taxa_file_snt = os.path.join(data_dir, 'taxa_subsets.json') | |
| # taxa_of_interest = datasets.get_taxa_of_interest(self.train_params['species_set'], self.train_params['num_aux_species'], | |
| # self.train_params['aux_species_seed'], self.train_params['taxa_file'], taxa_file_snt) | |
| obs_file = self.pos_eval_data_loc | |
| locs, labels, _, dates, _, _ = datasets.load_eval_inat_data(obs_file) | |
| unique_taxa, class_ids = np.unique(labels, return_inverse=True) | |
| class_to_taxa = unique_taxa.tolist() | |
| # idx_ss = datasets.get_idx_subsample_observations(labels, self.eval_params['num_samples'], random.randint(0,2**32), None, -1) | |
| idx_ss = datasets.get_idx_subsample_observations_eval(labels=labels, hard_cap=self.eval_params['num_samples']) | |
| locs = torch.from_numpy(np.array(locs)) | |
| labels = torch.from_numpy(np.array(class_ids)) | |
| locs = locs[idx_ss] | |
| labels = labels[idx_ss] | |
| with torch.no_grad(): | |
| pos_examples = {} | |
| for tt in self.taxa: | |
| c = class_to_taxa.index(tt) | |
| pos_examples[tt] = locs[labels == c] | |
| pos_examples[tt] = model(enc.encode(pos_examples[tt].to(self.eval_params['device'])), return_feats=True).cpu() | |
| # MAX VERSION # MAX VERSION # MAX VERSION | |
| # random negs | |
| neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| obs_file = self.background_eval_data_loc | |
| neg_locs, _, _, _, _, _ = datasets.load_eval_inat_data(obs_file) | |
| neg_locs = torch.from_numpy(neg_locs) | |
| if extra_input is not None: | |
| raise NotImplementedError('extra_input provided') | |
| # add target negs | |
| neg_examples = model(torch.cat([enc.encode(neg_examples, normalize=False), enc.encode( | |
| neg_locs[torch.randperm(neg_locs.shape[0], device=locs.device)[:10000]].clone().to( | |
| self.eval_params['device']), normalize=True)]), return_feats=True).cpu() | |
| loc_emb = model(loc_feat, return_feats=True) | |
| elif self.train_params['model'] == 'HyperNet': | |
| import datasets | |
| # from sklearn.linear_model import LogisticRegression | |
| # with open('paths.json', 'r') as f: | |
| # paths = json.load(f) | |
| # data_dir = paths['train'] | |
| # obs_file = os.path.join(data_dir, self.train_params['obs_file']) | |
| # taxa_file = os.path.join(data_dir, self.train_params['taxa_file']) | |
| # taxa_file_snt = os.path.join(data_dir, 'taxa_subsets.json') | |
| # | |
| # taxa_of_interest = datasets.get_taxa_of_interest(self.train_params['species_set'], self.train_params['num_aux_species'], | |
| # self.train_params['aux_species_seed'], self.train_params['taxa_file'], taxa_file_snt) | |
| # | |
| obs_file = self.pos_eval_data_loc | |
| locs, labels, _, dates, _, _ = datasets.load_eval_inat_data(obs_file) | |
| unique_taxa, class_ids = np.unique(labels, return_inverse=True) | |
| class_to_taxa = unique_taxa.tolist() | |
| if self.eval_params['num_samples'] > 0: | |
| # idx_ss = datasets.get_idx_subsample_observations(labels, self.eval_params['num_samples'], random.randint(0,2**32), None, -1) | |
| idx_ss = datasets.get_idx_subsample_observations_eval(labels=labels, hard_cap=self.eval_params['num_samples']) | |
| locs = torch.from_numpy(np.array(locs)[idx_ss]) | |
| labels = torch.from_numpy(np.array(class_ids)[idx_ss]) | |
| with torch.no_grad(): | |
| pos_examples = {} | |
| for tt in self.taxa: | |
| c = class_to_taxa.index(tt) | |
| pos_examples[tt] = locs[labels == c] | |
| pos_examples[tt] = model.pos_enc(enc.encode(pos_examples[tt].to(self.eval_params['device']))).cpu() | |
| # random negs | |
| neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| obs_file = self.background_eval_data_loc | |
| neg_locs, _, _, _, _, _ = datasets.load_eval_inat_data(obs_file) | |
| neg_locs = torch.from_numpy(neg_locs) | |
| if extra_input is not None: | |
| raise NotImplementedError('extra_input provided') | |
| neg_examples = model.pos_enc(torch.cat([enc.encode(neg_examples, normalize=False), enc.encode(neg_locs[torch.randperm(neg_locs.shape[0], device=locs.device)[:10000]].clone().to(self.eval_params['device']), normalize=True)])).cpu() | |
| loc_emb = model.pos_enc(loc_feat) | |
| #embs = torch.load(self.train_params['text_emb_path']) #TODO | |
| #embs1 = torch.load('experiments/gpt_data.pt', weights_only=False) | |
| embs1 = torch.load('experiments/gpt_data.pt', map_location='cpu') | |
| #embs1 = torch.load('ldsdm_data.pt') | |
| emb_ids1 = embs1['taxon_id'].tolist() | |
| keys1 = embs1['keys'] | |
| embs1 = embs1['data'] | |
| # embs2 doesn't even do anything. Could just remove the whole thing, but that is how it is in Max's code | |
| # MINE MINE MINE MINE MINE | |
| embs2 = torch.load('experiments/wiki_data_v4.pt') | |
| # MAX MAX MAX MAX | |
| # embs2 = torch.load('wiki_data_v3.pt') | |
| emb_ids2 = embs2['taxon_id'].tolist() | |
| keys2 = embs2['keys'] | |
| embs2 = embs2['data'] | |
| else: | |
| raise NotImplementedError('Eval for zero-shot not implemented') | |
| # if self.eval_params['num_samples'] == -1 and not (('CombinedModel' in self.train_params['model']) or ('MultiInputModel' in self.train_params['model'] or )): | |
| if self.eval_params['num_samples'] == -1 and not (self.train_params['model'] in ['CombinedModel', 'MultiInputModel', 'VariableInputModel', 'ResidualFCNet']): | |
| loc_emb = model.pos_enc(loc_feat) | |
| elif self.eval_params['num_samples'] == -1 and not (self.train_params['model'] in ['CombinedModel', 'MultiInputModel', 'VariableInputModel']): | |
| loc_emb = model.forward(loc_feat, return_feats=True) | |
| split_rng = np.random.default_rng(self.eval_params['split_seed']) | |
| write_gt_once = False | |
| #TODO: tt is the iNat taxa id for the taxa we are calculating AP for rn, tt_id is the index in the dictionary | |
| #ap_csv = "per_species_average_precision_valid.csv" | |
| #taxa_id_csv = "per_species_taxa_id_valid.csv" | |
| # with open(taxa_id_csv, mode='w', newline='') as csv_file: | |
| # writer = csv.writer(csv_file) | |
| # # If the array is multi-dimensional (e.g., 2D), iterate over rows | |
| # if isinstance(self.taxa, np.ndarray): | |
| # for value in self.taxa: | |
| # writer.writerow([value]) | |
| # else: | |
| # # If it's a flat array, directly write the values | |
| # writer.writerow(per_species_average_precision_valid) | |
| range, range_locs = [], [] | |
| for tt_id, tt in tqdm(enumerate(self.taxa)): | |
| class_of_interest = np.where(np.array(self.train_params['class_to_taxa']) == tt)[0] | |
| if len(class_of_interest) == 0 and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| # taxa of interest is not in the model | |
| results['per_species_average_precision_all'][tt_id] = np.nan | |
| # this only effects my models | |
| elif self.train_params['model'] == 'VariableInputModel': | |
| # generate ground truth labels for current taxa | |
| cur_loc_indices = np.array(self.loc_indices_per_species[tt_id]) | |
| cur_labels = np.array(self.labels_per_species[tt_id]) | |
| # apply per-species split: | |
| assert self.eval_params['split'] in ['all', 'val', 'test'] | |
| if self.eval_params['split'] != 'all': | |
| num_val = np.floor(len(cur_labels) * self.eval_params['val_frac']).astype(int) | |
| idx_rand = split_rng.permutation(len(cur_labels)) | |
| if self.eval_params['split'] == 'val': | |
| idx_sel = idx_rand[:num_val] | |
| elif self.eval_params['split'] == 'test': | |
| idx_sel = idx_rand[num_val:] | |
| cur_loc_indices = cur_loc_indices[idx_sel] | |
| cur_labels = cur_labels[idx_sel] | |
| cur_labels = (torch.from_numpy(cur_labels).to(self.eval_params['device']) > 0).float() | |
| with torch.no_grad(): | |
| logits = pred_mtx[:, tt_id] | |
| preds = torch.sigmoid(logits) | |
| #TODO metric value is calcuated | |
| #this is how we get the predictions, just matching the hexs for the spots we are interested in. | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer( | |
| cur_labels, | |
| preds[cur_loc_indices]).item() | |
| continue | |
| elif self.train_params['model'] == 'MultiInputModel': | |
| # generate ground truth labels for current taxa | |
| #todo: ask max, are the loc_indices the h3 indices at res 5? | |
| #these are the inidices of the locations of where we have evaluations | |
| cur_loc_indices = np.array(self.loc_indices_per_species[tt_id]) | |
| #loc_indices_per_species_array = np.array(self.loc_indices_per_species[tt_id]) | |
| #this is the answer key | |
| cur_labels = np.array(self.labels_per_species[tt_id]) #87373 "0." | |
| #labels_per_species_array = np.array(self.labels_per_species[tt_id]) #174746 '0' | |
| # apply per-species split: | |
| assert self.eval_params['split'] in ['all', 'val', 'test'] | |
| if self.eval_params['split'] != 'all': | |
| num_val = np.floor(len(cur_labels) * self.eval_params['val_frac']).astype(int) | |
| idx_rand = split_rng.permutation(len(cur_labels)) | |
| if self.eval_params['split'] == 'val': | |
| idx_sel = idx_rand[:num_val] | |
| elif self.eval_params['split'] == 'test': | |
| idx_sel = idx_rand[num_val:] | |
| cur_loc_indices = cur_loc_indices[idx_sel] | |
| cur_labels = cur_labels[idx_sel] | |
| cur_labels = (torch.from_numpy(cur_labels).to(self.eval_params['device']) > 0).float() | |
| #print('printing location testing') | |
| #matching_locations = obs_locs[loc_indices_per_species_array[labels_per_species_array == 1]]#21737 this is bigger because we take out the all and val locations | |
| matching_locations = obs_locs[cur_loc_indices[cur_labels == 1]] #10849 | |
| range_locs.append(matching_locations) | |
| #print(f'matching locations len: {len(matching_locations)}') | |
| range.append(cur_labels) | |
| #print(f'range cur labels len: {cur_labels.sum()}') | |
| #print(f'number of locations matches: matching locations: {np.shape(matching_locations)} and cur_labels: {cur_labels.sum()}') | |
| # if not write_gt_once: | |
| # snt_labels_csv = f"data/plot/taxa_locs/snt_locations_{tt}.csv" | |
| # with open(snt_labels_csv, mode='w', newline='') as csv_file: | |
| # writer = csv.writer(csv_file) | |
| # # If the array is multi-dimensional (e.g., 2D), iterate over rows | |
| # if isinstance(matching_locations, np.ndarray): | |
| # for value in matching_locations: | |
| # writer.writerow([value]) | |
| # else: | |
| # # If it's a flat array, directly write the values | |
| # writer.writerow(matching_locations) | |
| #print(f'current labels snt: {np.shape(cur_labels)}') | |
| with torch.no_grad(): | |
| logits = pred_mtx[:, tt_id] | |
| preds = torch.sigmoid(logits) | |
| #TODO metric value is calcuated | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer( | |
| cur_labels, | |
| preds[cur_loc_indices]).item() | |
| continue | |
| # MINE MINE MINE MINE MINE MINE | |
| # elif self.eval_params['num_samples'] == -1: | |
| # gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| # gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| # species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| # preds = loc_emb @ species_w.detach() | |
| # results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| # continue | |
| else: | |
| # generate ground truth labels for current taxa | |
| cur_loc_indices = np.array(self.loc_indices_per_species[tt_id]) | |
| cur_labels = np.array(self.labels_per_species[tt_id]) | |
| # apply per-species split: | |
| assert self.eval_params['split'] in ['all', 'val', 'test'] | |
| if self.eval_params['split'] != 'all': | |
| num_val = np.floor(len(cur_labels) * self.eval_params['val_frac']).astype(int) | |
| idx_rand = split_rng.permutation(len(cur_labels)) | |
| if self.eval_params['split'] == 'val': | |
| idx_sel = idx_rand[:num_val] | |
| elif self.eval_params['split'] == 'test': | |
| idx_sel = idx_rand[num_val:] | |
| cur_loc_indices = cur_loc_indices[idx_sel] | |
| cur_labels = cur_labels[idx_sel] | |
| cur_labels = (torch.from_numpy(cur_labels).to(self.eval_params['device']) > 0).float() | |
| ########################################################################################## | |
| # | |
| ########################################################################################## | |
| if self.eval_params['num_samples'] == -1 and self.train_params['model'] == 'HyperNet': | |
| species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| preds = loc_emb @ species_w.detach() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(cur_labels, | |
| preds[cur_loc_indices]).item() | |
| continue | |
| elif self.eval_params['num_samples'] == -1 and self.train_params['model'] == 'ResidualFCNet': | |
| preds = model.eval_single_class(x=loc_emb, class_of_interest=self.train_params['class_to_taxa'].index(tt)).detach() | |
| # species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| # preds = loc_emb @ species_w.detach() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(cur_labels, | |
| preds[cur_loc_indices]).item() | |
| continue | |
| if 'HyperNet' not in self.train_params['model'] and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| # extract model predictions for current taxa from prediction matrix | |
| pred = pred_mtx[cur_loc_indices, tt_id] | |
| elif self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0: | |
| if self.train_params['model'] == 'ResidualFCNet': | |
| if self.eval_params['num_samples'] == 0: | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| w = torch.nn.Parameter(torch.zeros(X.shape[1], 1, device=self.eval_params['device'])) | |
| nn.init.xavier_uniform_(w) | |
| pred = torch.sigmoid(((loc_emb @ w)))[cur_loc_indices].flatten() | |
| else: | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| y = torch.zeros(X.shape[0], dtype=torch.long, device=self.eval_params['device']) | |
| y[:pos_examples[tt].shape[0]] = 1 | |
| # MINE MINE MINE MINE MINE MINE | |
| # clf = LogisticRegression(class_weight='balanced', fit_intercept=False, C=0.05, max_iter=200, random_state=0).fit(X.numpy(), y.numpy()) | |
| # #pred = torch.from_numpy(clf.predict_proba(loc_emb.cpu()))[:,1] | |
| # pred = torch.sigmoid(((loc_emb @ (torch.from_numpy(clf.coef_).cuda().float().T)) + torch.from_numpy(clf.intercept_).cuda().float()).squeeze(-1))[cur_loc_indices] | |
| # MAX MAX MAX MAX MAX MAX MAX MAX MAX | |
| #clf = LogisticRegression(class_weight='balanced', fit_intercept=False, C=0.05, max_iter=200, random_state=0).fit(X.numpy(), y.numpy()) | |
| C = 0.05 | |
| w = torch.nn.Parameter(torch.zeros(X.shape[1], 1, device=self.eval_params['device'])) | |
| opt = torch.optim.Rprop([w], lr=0.001) | |
| crit = torch.nn.BCEWithLogitsLoss() | |
| crit2 = torch.nn.MSELoss() | |
| with torch.set_grad_enabled(True): | |
| for i in range(40): | |
| opt.zero_grad() | |
| output = X @ w | |
| yhat = y.float()[:, None] | |
| loss = 0.5 * crit(output[yhat == 0], yhat[yhat == 0]) + 0.5 * crit(output[yhat == 1], | |
| yhat[ | |
| yhat == 1]) + 1 / ( | |
| C * len(pos_examples[tt])) * crit2(w, 0 * w) | |
| loss.backward() | |
| opt.step() | |
| #pred = torch.from_numpy(clf.predict_proba(loc_emb.cpu()))[:,1] | |
| # pred = torch.sigmoid(((loc_emb @ w.cuda())))[cur_loc_indices].flatten() | |
| pred = torch.sigmoid(((loc_emb @ w)))[cur_loc_indices].flatten() | |
| #pred = torch.sigmoid(((loc_emb @ (torch.from_numpy(clf.coef_).cuda().float().T)) + torch.from_numpy(clf.intercept_).cuda().float()).squeeze(-1))[cur_loc_indices] | |
| elif self.train_params['model'] == 'HyperNet': | |
| if tt in emb_ids1: | |
| embs = embs1 | |
| emb_ids = emb_ids1 | |
| keys = keys1 | |
| else: | |
| print('yes') | |
| results['per_species_average_precision_all'][tt_id] = 0.0 | |
| continue | |
| embs = embs2 | |
| emb_ids = emb_ids2 | |
| keys = keys2 | |
| if tt not in emb_ids: | |
| results['per_species_average_precision_all'][tt_id] = 0.0 | |
| continue | |
| with torch.no_grad(): | |
| sec_ind = emb_ids.index(tt) | |
| sections = [i for i,x in enumerate(keys) if x[0] == sec_ind] | |
| def get_feat(x): | |
| species = model.species_enc(model.species_emb.zero_shot(x)) | |
| species_w, species_b = species[..., :-1], species[..., -1:] | |
| if self.eval_params['num_samples'] == 0: | |
| out = loc_emb @ (species_w.detach()).T | |
| return out | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| y = torch.zeros(X.shape[0], dtype=torch.long, device=self.eval_params['device']) | |
| y[:pos_examples[tt].shape[0]] = 1 | |
| C = 0.05 | |
| w = torch.nn.Parameter(torch.zeros_like(species_w,device=self.eval_params['device'])) | |
| opt = torch.optim.Rprop([w], lr=0.001) | |
| crit = torch.nn.BCEWithLogitsLoss() | |
| crit2 = torch.nn.MSELoss() | |
| with torch.set_grad_enabled(True): | |
| for i in range(40): | |
| opt.zero_grad() | |
| output = (X @ (w + species_w.detach()).T) + 0*species_b.squeeze(-1) | |
| yhat = y.float()[:, None].repeat(1, w.shape[0]) | |
| loss = 0.5*crit(output[yhat == 0], yhat[yhat == 0]) + 0.5*crit(output[yhat == 1], yhat[yhat == 1]) + \ | |
| 1/(C*len(pos_examples[tt])) * crit2(w, 0*w) | |
| loss.backward() | |
| opt.step() | |
| #print(i, loss.item()) | |
| #print(' ') | |
| out = loc_emb @ (w.data + species_w.detach()).T | |
| out = (out + 0*species_b.squeeze(-1)) | |
| return out | |
| # average precision score: | |
| yfeats = torch.cat([embs[section][None].to(self.eval_params['device']) for section in sections]) | |
| preds = get_feat(yfeats) | |
| if len(sections) > 1:#'habitat', 'overview_summary' | |
| kws = ['text', 'range', 'distribution', 'habitat'] if len(keys) == len(keys2) else [self.eval_params['text_section']] | |
| best_sections = [i for i,s in enumerate(sections) if any((x in keys[s][1].lower() for x in kws))] | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(cur_labels, preds[cur_loc_indices][:,best_sections].mean(dim=1)).item() | |
| else: | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(cur_labels, preds[cur_loc_indices][:,0].mean(dim=1)).item() | |
| continue | |
| else: | |
| raise NotImplementedError('Eval for hypernet not implemented') | |
| pred = preds[:,tt_id%32] | |
| # compute the AP for each taxa | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(cur_labels, pred).item() | |
| valid_taxa = ~np.isnan(results['per_species_average_precision_all']) | |
| # store results | |
| #TODO: this will have AP values for every species | |
| #tt_id | |
| per_species_average_precision_valid = results['per_species_average_precision_all'][valid_taxa] | |
| results['mean_average_precision'] = per_species_average_precision_valid.mean() | |
| results['num_eval_species_w_valid_ap'] = valid_taxa.sum() | |
| results['num_eval_species_total'] = len(self.taxa) | |
| taxas_and_ap_csv = "taxas_ap_range.csv" | |
| #ap_csv = "per_species_taxa_id_valid.csv" | |
| print(list(map(lambda row:len(row) ,range))) | |
| zipped_data = zip(self.taxa, per_species_average_precision_valid, list(map(lambda row:int(row.sum()),range)), range_locs) | |
| with open(taxas_and_ap_csv, mode='w', newline='') as csv_file: | |
| writer = csv.writer(csv_file) | |
| # Write the header (optional) | |
| writer.writerow(['Taxa ID', 'Average Precision','Range Size', 'Range']) | |
| # Write the zipped data | |
| for taxa, ap, range_size, tensor_range in zipped_data: | |
| # Flatten tensor to a single-line string | |
| tensor_range_str = format_tensor(tensor_range) | |
| writer.writerow([taxa, ap, range_size, tensor_range_str]) | |
| # with open(ap_csv, mode='w', newline='') as csv_file: | |
| # writer = csv.writer(csv_file) | |
| # # Write the zipped data | |
| # writer.writerows(per_species_average_precision_valid) | |
| return results | |
| def report(self, results): | |
| for field in ['mean_average_precision', 'num_eval_species_w_valid_ap', 'num_eval_species_total']: | |
| print(f'{field}: {results[field]}') | |
| class EvaluatorIUCN: | |
| def __init__(self, train_params, eval_params): | |
| self.train_params = train_params | |
| print(train_params['text_num_layers'],train_params['text_batchnorm'],train_params['text_hidden_dim'])#TODO | |
| self.eval_params = eval_params | |
| with open('paths.json', 'r') as f: | |
| paths = json.load(f) | |
| with open(os.path.join(paths['iucn'], 'iucn_res_5.json'), 'r') as f: | |
| self.data = json.load(f) | |
| self.obs_locs = np.array(self.data['locs'], dtype=np.float32) | |
| self.taxa = [int(tt) for tt in self.data['taxa_presence'].keys()] | |
| self.pos_eval_data_loc = os.path.join(paths['data'], 'positive_eval_data.npz') | |
| self.background_eval_data_loc = os.path.join(paths['data'], '10000_background_negs.npz') | |
| def run_evaluation(self, model, enc, extra_input=None): | |
| results = {} | |
| #self.train_params['model'] = 'ResidualFCNet' | |
| #m = model | |
| #model = lambda x, return_feats=True: m.pos_enc(x) | |
| results['per_species_average_precision_all'] = np.zeros(len(self.taxa), dtype=np.float32) | |
| # get eval locations and apply input encoding | |
| obs_locs = torch.from_numpy(self.obs_locs).to(self.eval_params['device']) | |
| loc_feat = torch.cat([enc.encode(obs_locs), extra_input.expand(obs_locs.shape[0], -1)], dim=1) if extra_input is not None else enc.encode(obs_locs) | |
| # get classes to eval | |
| # classes_of_interest = torch.zeros(len(self.taxa), dtype=torch.int64) | |
| classes_of_interest = np.zeros(len(self.taxa)) | |
| array_class_to_taxa = np.array(self.train_params['class_to_taxa']) | |
| for tt_id, tt in enumerate(self.taxa): | |
| class_of_interest = np.where(array_class_to_taxa == tt)[0] | |
| if len(class_of_interest) != 0: | |
| classes_of_interest[tt_id] = class_of_interest | |
| classes_of_interest = torch.from_numpy(classes_of_interest).to(dtype=torch.long, device=self.eval_params['device']) | |
| # MINE MINE MINE | |
| # classes_of_interest = classes_of_interest.to(self.eval_params['device']) | |
| if self.eval_params['extract_pos']: | |
| assert 'HyperNet' in self.train_params['model'] | |
| model = model.pos_enc | |
| self.train_params['model'] = 'ResidualFCNet' | |
| # Should only effect mine | |
| if ('CombinedModel' in self.train_params['model']) or ('MultiInputModel' in self.train_params['model']): | |
| with torch.no_grad(): | |
| dummy_context_mask = None | |
| dummy_context_sequence = None | |
| # generate model predictions for classes of interest at eval locations | |
| loc_emb = model(x=loc_feat, context_sequence=dummy_context_sequence, context_mask=dummy_context_mask, | |
| class_ids=classes_of_interest, return_feats=True) | |
| wt = model.get_eval_embeddings(classes_of_interest) | |
| print("Creating IUCN prediction matrix") | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif self.train_params['model'] == 'VariableInputModel': | |
| with torch.no_grad(): | |
| loc_emb = model.get_loc_emb(x=loc_feat) | |
| classes_of_interest = classes_of_interest.to(self.eval_params["device"]) | |
| wt = model.get_eval_embeddings(classes_of_interest) | |
| wt2 = model.get_ema_embeddings(classes_of_interest) | |
| # technically with my mock transformer I could just directly access the class embeddings but | |
| # I will need to use the emas when I move to the true transformer model (I think) | |
| # wt = model.class_emb.weight[classes_of_interest, :] | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif 'HyperNet' not in self.train_params['model'] and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| # generate model predictions for classes of interest at eval locations | |
| loc_emb = model(loc_feat, return_feats=True) | |
| wt = model.class_emb.weight[classes_of_interest, :] | |
| pred_mtx = torch.matmul(loc_emb, torch.transpose(wt, 0, 1)) | |
| elif (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| if self.train_params['model'] == 'ResidualFCNet': | |
| import datasets | |
| # from sklearn.linear_model import LogisticRegression | |
| # with open('paths.json', 'r') as f: | |
| # paths = json.load(f) | |
| # data_dir = paths['train'] | |
| # obs_file = os.path.join(data_dir, self.train_params['obs_file']) | |
| # taxa_file = os.path.join(data_dir, self.train_params['taxa_file']) | |
| # taxa_file_snt = os.path.join(data_dir, 'taxa_subsets.json') | |
| # | |
| # taxa_of_interest = datasets.get_taxa_of_interest(self.train_params['species_set'], self.train_params['num_aux_species'], | |
| # self.train_params['aux_species_seed'], self.train_params['taxa_file'], taxa_file_snt) | |
| obs_file = self.pos_eval_data_loc | |
| locs, labels, _, dates, _, _ = datasets.load_eval_inat_data(obs_file) | |
| unique_taxa, class_ids = np.unique(labels, return_inverse=True) | |
| class_to_taxa = unique_taxa.tolist() | |
| # idx_ss = datasets.get_idx_subsample_observations(labels, self.eval_params['num_samples'], random.randint(0,2**32), None, -1) | |
| idx_ss = datasets.get_idx_subsample_observations_eval(labels=labels, hard_cap=self.eval_params['num_samples']) | |
| locs = torch.from_numpy(np.array(locs)) | |
| labels = torch.from_numpy(np.array(class_ids)) | |
| locs = locs[idx_ss] | |
| labels = labels[idx_ss] | |
| # MINE MINE MINE MINE MINE MINE MINE MINE MINE | |
| # with torch.no_grad(): | |
| # pos_examples = {} | |
| # for tt in self.taxa: | |
| # c = class_to_taxa.index(tt) | |
| # pos_examples[tt] = locs[labels == c] | |
| # pos_examples[tt] = model(enc.encode(pos_examples[tt].to(self.eval_params['device'])), return_feats=True).cpu() | |
| # | |
| # if self.eval_params['target_background']: | |
| # target_background_dataset = datasets.get_train_data(params=self.train_params) | |
| # # print("CHECK IF THIS TARGET NEGS THING IS WORKING PROPERLY WHEN SERVER WORKS") | |
| # # print("IT MAY INCLUDE EVAL SPECIES / ONLY EVAL SPECIES") # it only includes the backbone species currently - good | |
| # | |
| # random_negs = utils.rand_samples(5000, self.eval_params['device'], rand_type='spherical') | |
| # | |
| # # Get the total number of locations | |
| # total_locs = len(target_background_dataset.locs) | |
| # | |
| # # If there are more than 5000 locations, sample 5000 | |
| # if total_locs > 5000: | |
| # indices = np.random.choice(total_locs, 5000, replace=False) | |
| # target_negs = target_background_dataset.locs[indices].to(self.eval_params['device']) | |
| # else: | |
| # target_negs = target_background_dataset.locs.to(self.eval_params['device']) | |
| # # print('CHECK THE FORMAT OF THESE TARGET LOCS COMPARED TO NEG LOCS') # look good | |
| # | |
| # neg_examples = torch.vstack((random_negs, target_negs)) | |
| # | |
| # del target_background_dataset | |
| # | |
| # else: | |
| # neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| # if extra_input is not None: | |
| # raise NotImplementedError('extra_input provided') | |
| # neg_examples = model(enc.encode(neg_examples, normalize=False), return_feats=True).cpu() | |
| # print("You can probably speed eval back up once the server is available by changing this shit back") | |
| # | |
| # # Function to process data in batches | |
| # def process_in_batches(model, loc_feat, batch_size=64): | |
| # loc_emb = [] | |
| # for i in range(0, len(loc_feat), batch_size): | |
| # batch = loc_feat[i:i + batch_size] | |
| # with torch.no_grad(): | |
| # batch_emb = model(batch, return_feats=True) | |
| # loc_emb.append(batch_emb) | |
| # return torch.cat(loc_emb, dim=0) # Concatenate the results | |
| # | |
| # # loc_emb = model(loc_feat, return_feats=True) | |
| # loc_emb = process_in_batches(model, loc_feat, batch_size=2048) | |
| pos_examples = {} | |
| for tt in self.taxa: | |
| c = class_to_taxa.index(tt) | |
| pos_examples[tt] = locs[labels == c] | |
| pos_examples[tt] = model(enc.encode(pos_examples[tt].to(self.eval_params['device'])), return_feats=True).cpu() | |
| obs_file = self.background_eval_data_loc | |
| neg_locs, _, _, _, _, _ = datasets.load_eval_inat_data(obs_file) | |
| neg_locs = torch.from_numpy(neg_locs) | |
| #random negs | |
| neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| if extra_input is not None: | |
| raise NotImplementedError('extra_input provided') | |
| # add target negs | |
| neg_examples = model(torch.cat([enc.encode(neg_examples, normalize=False), enc.encode(neg_locs[torch.randperm(neg_locs.shape[0], device=locs.device)[:10000]].clone().to(self.eval_params['device']), normalize=True)]), return_feats=True).cpu() | |
| loc_emb = model(loc_feat, return_feats=True) | |
| elif self.train_params['model'] == 'HyperNet': | |
| import datasets | |
| # from sklearn.linear_model import LogisticRegression | |
| # with open('paths.json', 'r') as f: | |
| # paths = json.load(f) | |
| # data_dir = paths['train'] | |
| # obs_file = os.path.join(data_dir, self.train_params['obs_file']) | |
| # taxa_file = os.path.join(data_dir, self.train_params['taxa_file']) | |
| # taxa_file_snt = os.path.join(data_dir, 'taxa_subsets.json') | |
| # | |
| # taxa_of_interest = datasets.get_taxa_of_interest(self.train_params['species_set'], self.train_params['num_aux_species'], | |
| # self.train_params['aux_species_seed'], self.train_params['taxa_file'], taxa_file_snt) | |
| obs_file = self.pos_eval_data_loc | |
| locs, labels, _, dates, _, _ = datasets.load_eval_inat_data(obs_file) | |
| # MINE MINE MINE MINE | |
| # unique_taxa, class_ids = np.unique(labels, return_inverse=True) | |
| # class_to_taxa = unique_taxa.tolist() | |
| # idx_ss = datasets.get_idx_subsample_observations(labels, self.eval_params['num_samples'], random.randint(0,2**32), None, -1) | |
| # locs = torch.from_numpy(np.array(locs)) | |
| # labels = torch.from_numpy(np.array(class_ids)) | |
| # locs = locs[idx_ss] | |
| # labels = labels[idx_ss] | |
| # with torch.no_grad(): | |
| # MAX MAX MAX MAX MAX MAX MAX | |
| unique_taxa, class_ids, class_counts = np.unique(labels, return_inverse=True, return_counts=True) | |
| class_counts = class_counts.clip(max=1000) | |
| if self.eval_params['num_samples'] > 0: | |
| class_to_taxa = unique_taxa.tolist() | |
| idx_ss = datasets.get_idx_subsample_observations_eval(labels=labels, hard_cap=self.eval_params['num_samples']) | |
| # idx_ss = datasets.get_idx_subsample_observations(labels, self.eval_params['num_samples'], random.randint(0,2**32), None, -1) | |
| locs = torch.from_numpy(np.array(locs)) | |
| labels = torch.from_numpy(np.array(class_ids)) | |
| locs = locs[idx_ss] | |
| labels = labels[idx_ss] | |
| pos_examples = {} | |
| for tt in self.taxa: | |
| c = class_to_taxa.index(tt) | |
| pos_examples[tt] = locs[labels == c] | |
| pos_examples[tt] = model.pos_enc(enc.encode(pos_examples[tt].to(self.eval_params['device']))).cpu() | |
| # MINE MINE MINE MINE MINE MINE MINE MINE | |
| # if self.eval_params['target_background']: | |
| # | |
| # target_background_dataset = datasets.get_train_data(params=self.train_params) | |
| # # print("CHECK IF THIS TARGET NEGS THING IS WORKING PROPERLY WHEN SERVER WORKS") | |
| # # print("IT MAY INCLUDE EVAL SPECIES / ONLY EVAL SPECIES") | |
| # | |
| # random_negs = utils.rand_samples(5000, self.eval_params['device'], rand_type='spherical') | |
| # | |
| # # Get the total number of locations | |
| # total_locs = len(target_background_dataset.locs) | |
| # | |
| # # If there are more than 5000 locations, sample 5000 | |
| # if total_locs > 5000: | |
| # indices = np.random.choice(total_locs, 5000, replace=False) | |
| # target_negs = target_background_dataset.locs[indices].to(self.eval_params['device']) | |
| # else: | |
| # target_negs = target_background_dataset.locs.to(self.eval_params['device']) | |
| # # print('CHECK THE FORMAT OF THESE TARGET LOCS COMPARED TO NEG LOCS') | |
| # | |
| # neg_examples = torch.vstack((random_negs, target_negs)) | |
| # | |
| # del target_background_dataset | |
| # | |
| # else: | |
| # neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| # MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX | |
| obs_file = self.background_eval_data_loc | |
| neg_locs, _, _, _, _, _ = datasets.load_eval_inat_data(obs_file) | |
| neg_locs = torch.from_numpy(neg_locs) | |
| # random negs | |
| neg_examples = utils.rand_samples(10000, self.eval_params['device'], rand_type='spherical') | |
| if extra_input is not None: | |
| raise NotImplementedError('extra_input provided') | |
| # MINE MINE MINE | |
| # neg_examples = model.pos_enc(enc.encode(neg_examples, normalize=False)).cpu() | |
| # MAX MAX MAX MAX MAX MAX MAX | |
| # add target negs | |
| neg_examples = model.pos_enc(torch.cat([enc.encode(neg_examples, normalize=False), enc.encode(neg_locs[torch.randperm(neg_locs.shape[0], device=locs.device)[:10000]].clone().to(self.eval_params['device']), normalize=True)])).cpu() | |
| #embs = torch.load(self.train_params['text_emb_path']) #TODO | |
| embs = torch.load('gpt_data.pt', weights_only=False) | |
| #embs = torch.load('ldsdm_data.pt') | |
| emb_ids = embs['taxon_id'].tolist() | |
| keys = embs['keys'] | |
| embs = embs['data'] | |
| # embs2 doesn't even do anything. Could just remove the whole thing, but that is how it is in Max's code | |
| # MINE MINE MINE | |
| embs2 = torch.load('wiki_data_v4.pt', weights_only=False) | |
| # MAX MAX MAX | |
| # embs2 = torch.load('wiki_data_v3.pt') | |
| emb_ids2 = embs2['taxon_id'].tolist() | |
| keys2 = embs2['keys'] | |
| embs2 = embs2['data'] | |
| loc_emb = model.pos_enc(loc_feat) | |
| else: | |
| raise NotImplementedError('Eval for zero-shot not implemented') | |
| # MINE - my version - why am I stopping residualFCnets doing this? | |
| # if self.eval_params['num_samples'] == -1 and not (('CombinedModel' in self.train_params['model']) or ('MultiInputModel' in self.train_params['model']) or ('ResidualFCNet' in self.train_params['model'])): | |
| # MAX - a variant of Maxs - only difference should now be my model types | |
| #if self.eval_params['num_samples'] == -1 and not (('CombinedModel' in self.train_params['model']) or ('MultiInputModel' in self.train_params['model'])): | |
| if self.eval_params['num_samples'] == -1 and not (self.train_params['model'] in ['CombinedModel', 'MultiInputModel', 'VariableInputModel', 'ResidualFCNet']): | |
| loc_emb = model.pos_enc(loc_feat) | |
| if self.eval_params['num_samples'] == -1 and not (self.train_params['model'] in ['CombinedModel', 'MultiInputModel', 'VariableInputModel']): | |
| loc_emb = model.forward(loc_feat, return_feats=True) | |
| for tt_id, tt in tqdm(enumerate(self.taxa)): | |
| class_of_interest = np.where(array_class_to_taxa == tt)[0] | |
| if len(class_of_interest) == 0 and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| # taxa of interest is not in the model | |
| results['per_species_average_precision_all'][tt_id] = np.nan | |
| else: | |
| # Only effects my models | |
| if self.train_params['model'] == 'MultiInputModel': | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| with torch.no_grad(): | |
| logits = pred_mtx[:, tt_id] | |
| preds = torch.sigmoid(logits) | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| continue | |
| elif self.train_params['model'] == 'VariableInputModel': | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| with torch.no_grad(): | |
| logits = pred_mtx[:, tt_id] | |
| preds = torch.sigmoid(logits) | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| continue | |
| # MINE MINE MINE | |
| # elif (self.train_params['model'] == 'ResidualFCNet') and (self.eval_params['num_samples'] <= 0): | |
| # gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| # gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| # with torch.no_grad(): | |
| # logits = pred_mtx[:, tt_id] | |
| # preds = torch.sigmoid(logits) | |
| # results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| # continue | |
| if self.eval_params['num_samples'] == -1 and self.train_params['model'] == 'HyperNet': | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| preds = loc_emb @ species_w.detach() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| continue | |
| elif self.eval_params['num_samples'] == -1 and self.train_params['model'] == 'ResidualFCNet': | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| preds = model.eval_single_class(x=loc_emb, class_of_interest=self.train_params['class_to_taxa'].index(tt)).detach() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| continue | |
| # MINE MINE MINE MINE MINE MINE MINE - seems un needed? | |
| # elif (self.eval_params['num_samples'] == -1) and ('Hypernet' in self.train_params['model']): | |
| # gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| # gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| # species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| # preds = loc_emb @ species_w.detach() | |
| # results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds).item() | |
| # continue | |
| # extract model predictions for current taxa from prediction matrix | |
| if 'HyperNet' not in self.train_params['model'] and not (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| pred = pred_mtx[:, tt_id] | |
| elif (self.train_params['zero_shot'] or self.eval_params['num_samples'] > 0): | |
| if self.train_params['model'] == 'ResidualFCNet': | |
| if self.eval_params['num_samples'] == 0: | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| w = torch.nn.Parameter(torch.zeros(X.shape[1], 1, device=self.eval_params['device'])) | |
| nn.init.xavier_uniform_(w) | |
| pred = torch.sigmoid(((loc_emb @ w)))[cur_loc_indices].flatten() | |
| else: | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| y = torch.zeros(X.shape[0], dtype=torch.long, device=self.eval_params['device']) | |
| y[:pos_examples[tt].shape[0]] = 1 | |
| # MINE MINE MINE | |
| # clf = LogisticRegression(class_weight='balanced', fit_intercept=False, C=0.05, max_iter=200, random_state=0).fit(X.numpy(), y.numpy()) | |
| # #pred = torch.from_numpy(clf.predict_proba(loc_emb.cpu()))[:,1] | |
| # pred = torch.sigmoid(((loc_emb @ (torch.from_numpy(clf.coef_).to(self.eval_params['device']).float().T)) + torch.from_numpy(clf.intercept_).to(self.eval_params['device']).float()).squeeze(-1)) | |
| # # pred = torch.sigmoid(((loc_emb @ (torch.from_numpy(clf.coef_).cuda().float().T)) + torch.from_numpy(clf.intercept_).cuda().float()).squeeze(-1)) | |
| # MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX MAX | |
| #clf = LogisticRegression(class_weight='balanced', fit_intercept=False, C=0.05, max_iter=200, random_state=0).fit(X.numpy(), y.numpy()) | |
| C = 0.05 | |
| w = torch.nn.Parameter(torch.zeros(X.shape[1], 1, device=self.eval_params['device'])) | |
| opt = torch.optim.Rprop([w], lr=0.001) | |
| crit = torch.nn.BCEWithLogitsLoss() | |
| crit2 = torch.nn.MSELoss() | |
| with torch.set_grad_enabled(True): | |
| for i in range(40): | |
| opt.zero_grad() | |
| output = X @ w | |
| yhat = y.float()[:, None] | |
| loss = 0.5 * crit(output[yhat == 0], yhat[yhat == 0]) + 0.5 * crit(output[yhat == 1], | |
| yhat[ | |
| yhat == 1]) + 1 / ( | |
| C * len(pos_examples[tt])) * crit2(w, 0 * w) | |
| loss.backward() | |
| opt.step() | |
| pred = torch.sigmoid(((loc_emb @ w))).flatten() | |
| #pred = torch.from_numpy(clf.predict_proba(loc_emb.cpu()))[:,1] | |
| #pred = torch.sigmoid(((loc_emb @ (torch.from_numpy(clf.coef_).cuda().float().T)) + torch.from_numpy(clf.intercept_).cuda().float()).squeeze(-1)) | |
| #locs = torch.from_numpy(utils.coord_grid((1000,2000))).to(self.eval_params['device']) | |
| #locs = model(enc.encode(locs), return_feats=True) | |
| #img = torch.sigmoid(((locs @ (torch.from_numpy(clf.coef_).cuda().float().T)) + torch.from_numpy(clf.intercept_).cuda().float()).squeeze(-1)) | |
| #plt.imshow(img.detach().cpu()) | |
| elif self.train_params['model'] == 'HyperNet': | |
| if tt not in emb_ids and tt not in emb_ids2: | |
| results['per_species_average_precision_all'][tt_id] = 0.0 | |
| continue | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| if self.eval_params['num_samples'] == -1: | |
| species_w = model.species_params[self.train_params['class_to_taxa'].index(tt)] | |
| preds = loc_emb @ species_w.detach() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt,preds).item() | |
| continue | |
| with torch.no_grad(): | |
| if tt in emb_ids: | |
| em = embs | |
| emi = emb_ids | |
| ky = keys | |
| else: | |
| results['per_species_average_precision_all'][tt_id] = 0.0 | |
| continue | |
| em = embs2 | |
| emi = emb_ids2 | |
| ky = keys2 | |
| sec_ind = emi.index(tt) | |
| sections = [i for i,x in enumerate(ky) if x[0] == sec_ind] | |
| order = ['distribution', 'range', 'text'] | |
| best_section = None | |
| order_ind = 0 | |
| while best_section is None and order_ind < len(order): | |
| for section in sections: | |
| if order[order_ind] in ky[section][1].lower(): | |
| best_section = section | |
| break | |
| order_ind += 1 | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| def get_feat(x): | |
| species = model.species_enc(model.species_emb.zero_shot(x)) | |
| species_w, species_b = species[..., :-1], species[..., -1:] | |
| if self.eval_params['num_samples'] == 0: | |
| out = loc_emb @ (species_w.detach()).T | |
| return out | |
| X = torch.cat([pos_examples[tt], neg_examples], dim=0).to(self.eval_params['device']) | |
| y = torch.zeros(X.shape[0], dtype=torch.long, device=self.eval_params['device']) | |
| y[:pos_examples[tt].shape[0]] = 1 | |
| C = 0.05 | |
| w = torch.nn.Parameter(torch.zeros_like(species_w, device=self.eval_params['device'])) | |
| opt = torch.optim.Rprop([w], lr=0.001) | |
| crit = torch.nn.BCEWithLogitsLoss() | |
| crit2 = torch.nn.MSELoss() | |
| with torch.set_grad_enabled(True): | |
| for i in range(40): | |
| opt.zero_grad() | |
| output = (X @ (w + species_w.detach()).T) + 0*species_b.squeeze(-1) | |
| yhat = y.float()[:, None].repeat(1, w.shape[0]) | |
| loss = 0.5 * crit(output[yhat == 0], yhat[yhat == 0]) + 0.5 * crit( | |
| output[yhat == 1], yhat[yhat == 1]) + 1 / ( | |
| C * len(pos_examples[tt])) * crit2(w, 0 * w) | |
| loss.backward() | |
| opt.step() | |
| '''out = loc_emb @ (w.data + species_w.detach()).T | |
| gt = torch.zeros(out.shape[0], dtype=torch.float32, | |
| device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| print(utils.average_precision_score_fasterer(gt, out[:, 0]).item())''' | |
| out = loc_emb @ (w.data + species_w.detach()).T | |
| out = (out + 0*species_b.squeeze(-1)) | |
| return out | |
| # average precision score: | |
| yfeats = torch.cat([em[section][None].to(self.eval_params['device']) for section in sections]) | |
| preds = get_feat(yfeats) | |
| if len(sections) > 1:#'habitat', 'overview_summary' | |
| kws = [self.eval_params['text_section']] if len(ky) == len(keys) else ['text', 'range','distribution','habitat'] | |
| best_sections = [i for i,s in enumerate(sections) if any((x in ky[s][1].lower() for x in kws))] | |
| #yfeats2 = torch.cat( | |
| # [em[section][None].to(self.eval_params['device']) for section in best_sections]).mean(dim=0, keepdim=True) | |
| #pred2 = get_feat(yfeats2) | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds[:, best_sections].mean(dim=1)).item() | |
| else: | |
| # MINE MINE MINE MINE | |
| # sigmoid_preds = torch.sigmoid(preds[:, 0]) | |
| # results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, sigmoid_preds).item() | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, preds[:, 0]).item() | |
| continue | |
| else: | |
| if tt_id % 32 == 0: | |
| # MINE MINE MINE MINE | |
| # with torch.no_grad(): | |
| # preds = torch.empty(loc_feat.shape[0], classes_of_interest[tt_id:tt_id+32].shape[0], device=self.eval_params['device']) | |
| # for i in range(0,preds.shape[0],50000): | |
| # xbatch = loc_feat[i:i+50000] | |
| # ybatch = classes_of_interest[tt_id:tt_id+32].to(self.eval_params['device']).expand(xbatch.shape[0], -1) | |
| # preds[i:i+50000] = model(xbatch, ybatch) | |
| preds = torch.empty(loc_feat.shape[0], classes_of_interest[tt_id:tt_id+32].shape[0], device=self.eval_params['device']) | |
| for i in range(0,preds.shape[0],50000): | |
| xbatch = loc_feat[i:i+50000] | |
| ybatch = classes_of_interest[tt_id:tt_id+32].to(self.eval_params['device']).expand(xbatch.shape[0], -1) | |
| preds[i:i+50000] = model(xbatch, ybatch) | |
| pred = preds[:,tt_id%32] | |
| gt = torch.zeros(obs_locs.shape[0], dtype=torch.float32, device=self.eval_params['device']) | |
| gt[self.data['taxa_presence'][str(tt)]] = 1.0 | |
| # average precision score: | |
| results['per_species_average_precision_all'][tt_id] = utils.average_precision_score_fasterer(gt, pred).item() | |
| valid_taxa = ~np.isnan(results['per_species_average_precision_all']) | |
| # store results | |
| per_species_average_precision_valid = results['per_species_average_precision_all'][valid_taxa] | |
| results['mean_average_precision'] = per_species_average_precision_valid.mean() | |
| results['num_eval_species_w_valid_ap'] = valid_taxa.sum() | |
| results['num_eval_species_total'] = len(self.taxa) | |
| return results | |
| def report(self, results): | |
| for field in ['mean_average_precision', 'num_eval_species_w_valid_ap', 'num_eval_species_total']: | |
| print(f'{field}: {results[field]}') | |
| # MINE MINE MINE MINE but shouldn't effect things too much | |
| def batched_matmul(self, loc_emb, wt): | |
| batch_size = self.eval_params["batch_size"] | |
| num_samples = loc_emb.size(0) | |
| num_batches = (num_samples + batch_size - 1) // batch_size # Ensures rounding up | |
| # Preallocate the result array | |
| pred_mtx = np.empty((num_samples, wt.size(0)), dtype=np.float32) | |
| wt_T = wt.t() | |
| # Buffer size for temporary storage | |
| buffer_size = batch_size * 10 # Adjust buffer size as needed | |
| buffer = np.empty((buffer_size, wt.size(0)), dtype=np.float32) | |
| buffer_index = 0 | |
| current_write_index = 0 | |
| for _, i in tqdm(enumerate(range(num_batches))): | |
| start_idx = i * batch_size | |
| end_idx = min(start_idx + batch_size, num_samples) | |
| # Perform matrix multiplication for the current batch in PyTorch | |
| loc_emb_batch = loc_emb[start_idx:end_idx].to(self.eval_params['device']) | |
| batch_result = torch.matmul(loc_emb_batch, wt_T).cpu().numpy() | |
| # Calculate the size of the current batch | |
| current_batch_size = end_idx - start_idx | |
| # Check if the buffer can accommodate the current batch | |
| if buffer_index + current_batch_size > buffer_size: | |
| # Write buffer contents to pred_mtx | |
| pred_mtx[current_write_index:current_write_index + buffer_index] = buffer[:buffer_index] | |
| current_write_index += buffer_index | |
| buffer_index = 0 # Reset buffer index | |
| # Add the current batch result to the buffer | |
| buffer[buffer_index:buffer_index + current_batch_size] = batch_result | |
| buffer_index += current_batch_size | |
| # Clean up to free memory | |
| del loc_emb_batch | |
| del batch_result | |
| # torch.cuda.empty_cache() # Consider removing if unnecessary | |
| # Write any remaining data in the buffer to pred_mtx | |
| if buffer_index > 0: | |
| pred_mtx[current_write_index:current_write_index + buffer_index] = buffer[:buffer_index] | |
| return pred_mtx | |
| class EvaluatorGeoPrior: | |
| def __init__(self, train_params, eval_params): | |
| # store parameters: | |
| self.train_params = train_params | |
| self.eval_params = eval_params | |
| with open('paths.json', 'r') as f: | |
| paths = json.load(f) | |
| # load vision model predictions: | |
| self.data = np.load(os.path.join(paths['geo_prior'], 'geo_prior_model_preds.npz')) | |
| print(self.data['probs'].shape[0], 'total test observations') | |
| # load locations: | |
| meta = pd.read_csv(os.path.join(paths['geo_prior'], 'geo_prior_model_meta.csv')) | |
| self.obs_locs = np.vstack((meta['longitude'].values, meta['latitude'].values)).T.astype(np.float32) | |
| temp = np.array(meta['observed_on'].values, dtype='S10') | |
| temp = temp.view('S1').reshape((temp.size, -1)) | |
| years = temp[:, :4].view('S4').astype(int)[:, 0] | |
| months = temp[:, 5:7].view('S2').astype(int)[:, 0] | |
| days = temp[:, 8:10].view('S2').astype(int)[:, 0] | |
| days_per_month = np.cumsum([0] + [monthrange(2018, mm)[1] for mm in range(1, 12)]) | |
| dates = days_per_month[months - 1] + days - 1 | |
| self.dates = np.round((dates) / 365.0, 4).astype(np.float32) | |
| # taxonomic mapping: | |
| self.taxon_map = self.find_mapping_between_models(self.data['model_to_taxa'], self.train_params['class_to_taxa']) | |
| self.time_enc = utils.TimeEncoder() if train_params['input_time'] else None | |
| print(self.taxon_map.shape[0], 'out of', len(self.data['model_to_taxa']), 'taxa in both vision and geo models') | |
| cs = torch.load('class_counts.pt') | |
| cs = cs.sum() / cs | |
| cs = cs.to(self.eval_params['device']) | |
| self.C = cs[None] | |
| self.pdf = utils.DataPDFH3(device=self.eval_params['device']) | |
| def find_mapping_between_models(self, vision_taxa, geo_taxa): | |
| # this will output an array of size N_overlap X 2 | |
| # the first column will be the indices of the vision model, and the second is their | |
| # corresponding index in the geo model | |
| taxon_map = np.ones((vision_taxa.shape[0], 2), dtype=np.int32)*-1 | |
| taxon_map[:, 0] = np.arange(vision_taxa.shape[0]) | |
| geo_taxa_arr = np.array(geo_taxa) | |
| for tt_id, tt in enumerate(vision_taxa): | |
| ind = np.where(geo_taxa_arr==tt)[0] | |
| if len(ind) > 0: | |
| taxon_map[tt_id, 1] = ind[0] | |
| inds = np.where(taxon_map[:, 1]>-1)[0] | |
| taxon_map = taxon_map[inds, :] | |
| return taxon_map | |
| def convert_to_inat_vision_order(self, geo_pred_ip, vision_top_k_prob, vision_top_k_inds, vision_taxa, taxon_map, k=1.0): | |
| # this is slow as we turn the sparse input back into the same size as the dense one | |
| vision_pred = np.zeros((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32) | |
| geo_pred = k*np.ones((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32) | |
| vision_pred[np.arange(vision_pred.shape[0])[..., np.newaxis], vision_top_k_inds] = vision_top_k_prob | |
| geo_pred[:, taxon_map[:, 0]] = geo_pred_ip[:, taxon_map[:, 1]] | |
| return geo_pred, vision_pred | |
| def run_evaluation(self, model, enc, extra_input=None): | |
| results = {} | |
| # loop over in batches | |
| batch_start = np.hstack((np.arange(0, self.data['probs'].shape[0], self.eval_params['batch_size']), self.data['probs'].shape[0])) | |
| correct_pred = np.zeros(self.data['probs'].shape[0]) | |
| from tqdm import tqdm | |
| for bb_id, bb in tqdm(enumerate(range(len(batch_start)-1))): | |
| batch_inds = np.arange(batch_start[bb], batch_start[bb+1]) | |
| vision_probs = self.data['probs'][batch_inds, :] | |
| vision_inds = self.data['inds'][batch_inds, :] | |
| gt = self.data['labels'][batch_inds] | |
| dates = torch.from_numpy(self.dates[batch_inds]) | |
| obs_locs_batch = torch.from_numpy(self.obs_locs[batch_inds, :]).to(self.eval_params['device']) | |
| noise_level = 1.0 | |
| if self.time_enc is not None: | |
| extra_input = self.time_enc.encode(torch.cat([dates[...,None], torch.full((*dates.shape, 1),noise_level)], dim=1)).to( | |
| self.eval_params['device']) | |
| loc_feat = torch.cat([enc.encode(obs_locs_batch), extra_input], 1) if extra_input is not None else enc.encode(obs_locs_batch) | |
| with torch.no_grad(): | |
| geo_pred = model(loc_feat).cpu().numpy() | |
| geo_pred, vision_pred = self.convert_to_inat_vision_order(geo_pred, vision_probs, vision_inds, | |
| self.data['model_to_taxa'], self.taxon_map, k=1.0) | |
| #geo_pred = softmax(torch.from_numpy(geo_pred), dim=1).numpy() | |
| comb_pred = np.argmax(vision_pred*geo_pred, 1) | |
| comb_pred = (comb_pred==gt) | |
| correct_pred[batch_inds] = comb_pred | |
| accuracy_by_taxa = np.zeros(len(self.data['model_to_taxa'])) | |
| for tt_id, tt in enumerate(self.data['model_to_taxa']): | |
| inds = np.where(self.data['labels'] == tt)[0] | |
| accuracy_by_taxa[tt_id] = float((correct_pred[inds].mean())) | |
| torch.save(correct_pred, f'correct_{noise_level}.pt') | |
| torch.save(accuracy_by_taxa, f'abt_{noise_level}.pt') | |
| results['vision_only_top_1'] = float((self.data['inds'][:, -1] == self.data['labels']).mean()) | |
| results['vision_geo_top_1'] = float(correct_pred.mean()) | |
| return results | |
| def report(self, results): | |
| print('Overall accuracy vision only model', round(results['vision_only_top_1'], 3)) | |
| print('Overall accuracy of geo model ', round(results['vision_geo_top_1'], 3)) | |
| print('Gain ', round(results['vision_geo_top_1'] - results['vision_only_top_1'], 3)) | |
| class EvaluatorGeoFeature: | |
| def __init__(self, train_params, eval_params): | |
| self.train_params = train_params | |
| self.eval_params = eval_params | |
| with open('paths.json', 'r') as f: | |
| paths = json.load(f) | |
| self.data_path = paths['geo_feature'] | |
| self.country_mask = tifffile.imread(os.path.join(paths['masks'], 'USA_MASK.tif')) == 1 | |
| self.raster_names = ['ABOVE_GROUND_CARBON', 'ELEVATION', 'LEAF_AREA_INDEX', 'NON_TREE_VEGITATED', 'NOT_VEGITATED', 'POPULATION_DENSITY', 'SNOW_COVER', 'SOIL_MOISTURE', 'TREE_COVER'] | |
| self.raster_names_log_transform = ['POPULATION_DENSITY'] | |
| def load_raster(self, raster_name, log_transform=False): | |
| raster = tifffile.imread(os.path.join(self.data_path, raster_name + '.tif')).astype(np.float32) | |
| valid_mask = ~np.isnan(raster).copy() & self.country_mask | |
| # log scaling: | |
| if log_transform: | |
| raster[valid_mask] = np.log1p(raster[valid_mask] - raster[valid_mask].min()) | |
| # 0/1 scaling: | |
| raster[valid_mask] -= raster[valid_mask].min() | |
| raster[valid_mask] /= raster[valid_mask].max() | |
| return raster, valid_mask | |
| def get_split_labels(self, raster, split_ids, split_of_interest): | |
| # get the GT labels for a subset | |
| inds_y, inds_x = np.where(split_ids==split_of_interest) | |
| return raster[inds_y, inds_x] | |
| def get_split_feats(self, model, enc, split_ids, split_of_interest, extra_input=None): | |
| locs = utils.coord_grid(self.country_mask.shape, split_ids=split_ids, split_of_interest=split_of_interest) | |
| locs = torch.from_numpy(locs).to(self.eval_params['device']) | |
| locs_enc = torch.cat([enc.encode(locs), extra_input.expand(locs.shape[0], -1)], 1) if extra_input is not None else enc.encode(locs) | |
| with torch.no_grad(): | |
| feats = model(locs_enc, return_feats=True).cpu().numpy() | |
| return feats | |
| def run_evaluation(self, model2, enc, extra_input=None): | |
| if self.train_params['model'] == 'ResidualFCNet': | |
| model = model2 | |
| elif self.train_params['model'] == 'HyperNet': | |
| model = lambda x, return_feats=True: model2.pos_enc(x) | |
| else: | |
| raise NotImplementedError() | |
| results = {} | |
| for raster_name in self.raster_names: | |
| do_log_transform = raster_name in self.raster_names_log_transform | |
| raster, valid_mask = self.load_raster(raster_name, do_log_transform) | |
| split_ids = utils.create_spatial_split(raster, valid_mask, cell_size=self.eval_params['cell_size']) | |
| feats_train = self.get_split_feats(model, enc, split_ids=split_ids, split_of_interest=1, extra_input=extra_input) | |
| feats_test = self.get_split_feats(model, enc, split_ids=split_ids, split_of_interest=2, extra_input=extra_input) | |
| labels_train = self.get_split_labels(raster, split_ids, 1) | |
| labels_test = self.get_split_labels(raster, split_ids, 2) | |
| scaler = MinMaxScaler() | |
| feats_train_scaled = scaler.fit_transform(feats_train) | |
| feats_test_scaled = scaler.transform(feats_test) | |
| clf = RidgeCV(alphas=(0.1, 1.0, 10.0), cv=10, fit_intercept=True, scoring='r2').fit(feats_train_scaled, labels_train) | |
| train_score = clf.score(feats_train_scaled, labels_train) | |
| test_score = clf.score(feats_test_scaled, labels_test) | |
| results[f'train_r2_{raster_name}'] = float(train_score) | |
| results[f'test_r2_{raster_name}'] = float(test_score) | |
| results[f'alpha_{raster_name}'] = float(clf.alpha_) | |
| return results | |
| def report(self, results): | |
| report_fields = [x for x in results if 'test_r2' in x] | |
| for field in report_fields: | |
| print(f'{field}: {results[field]}') | |
| print(np.mean([results[field] for field in report_fields])) | |
| # I need train overrides for some of my stuff but it should have zero impact on other things | |
| def launch_eval_run(overrides, train_overrides=None): | |
| eval_params = setup.get_default_params_eval(overrides) | |
| # set up model: | |
| eval_params['model_path'] = os.path.join(eval_params['exp_base'], eval_params['experiment_name'], eval_params['ckp_name']) | |
| #train_params = torch.load(eval_params['model_path'], map_location='cpu', weights_only=False) | |
| train_params = torch.load(eval_params['model_path'], map_location='cpu') | |
| default_params = setup.get_default_params_train() | |
| for key in default_params: | |
| if key not in train_params['params']: | |
| train_params['params'][key] = default_params[key] | |
| # MINE - this is hopefully just for my models - must ensure this - should have zero impact on hypernets | |
| if train_overrides != None: | |
| for key, value in train_overrides.items(): | |
| #print(f'updating train param {key}') | |
| train_params['params'][key] = value | |
| model = models.get_model(train_params['params'], inference_only=True) | |
| model.load_state_dict(train_params['state_dict'], strict=False) | |
| model = model.to(eval_params['device']) | |
| model.eval() | |
| # create input encoder: | |
| if train_params['params']['input_enc'] in ['env', 'sin_cos_env', 'sh_env']: | |
| raster = datasets.load_env().to(eval_params['device']) | |
| else: | |
| raster = None | |
| enc = utils.CoordEncoder(train_params['params']['input_enc'], raster=raster, input_dim=train_params['params']['input_dim']) | |
| if train_params['params']['input_time']: | |
| time_enc = utils.TimeEncoder(input_enc='conical') if train_params['params']['input_time'] else None | |
| extra_input = torch.cat([time_enc.encode(torch.tensor([[0.0, 1.0]]))], dim=1).to(eval_params['device']) | |
| else: | |
| extra_input = None | |
| # This should only effect my models | |
| # This is where I create the eval "species tokens" from the specified number of context points | |
| # TODO just use the existing train params and some if statements to get the right dataset without having to use train overides | |
| if train_params['params']['model'] == 'MultiInputModel': | |
| train_dataset = datasets.get_train_data(train_params['params']) | |
| if 'text' in train_params['params']['dataset']: | |
| if eval_params['text_section'] != '': | |
| train_dataset.select_text_section(eval_params['text_section']) | |
| print(f'Using {eval_params["text_section"]} text for evaluation') | |
| train_loader = torch.utils.data.DataLoader( | |
| train_dataset, | |
| batch_size=train_params['params']['batch_size'], | |
| shuffle=True, | |
| num_workers=8, | |
| collate_fn=getattr(train_dataset, 'collate_fn', None)) | |
| # if len(train_params['params']['class_to_taxa']) != train_dataset.class_to_taxa: | |
| # Create new embedding layers for the expanded classes | |
| num_new_classes = len(train_dataset.class_to_taxa) | |
| embedding_dim = model.ema_embeddings.embedding_dim | |
| new_ema_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| new_eval_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| nn.init.xavier_uniform_(new_ema_embeddings.weight) | |
| nn.init.xavier_uniform_(new_eval_embeddings.weight) | |
| # Convert lists to numpy arrays for indexing | |
| class_to_taxa_np = np.array(train_params['params']['class_to_taxa']) | |
| class_to_taxa_expanded_np = np.array(train_dataset.class_to_taxa) | |
| # Find common taxa and their indices | |
| common_taxa, original_indices, expanded_indices = np.intersect1d( | |
| class_to_taxa_np, class_to_taxa_expanded_np, return_indices=True) | |
| # Update new embeddings for the common taxa | |
| new_ema_embeddings.weight.data[expanded_indices] = model.ema_embeddings.weight.data[original_indices] | |
| new_eval_embeddings.weight.data[expanded_indices] = model.eval_embeddings.weight.data[original_indices] | |
| # Replace old embeddings with new embeddings | |
| model.ema_embeddings = new_ema_embeddings | |
| model.eval_embeddings = new_eval_embeddings | |
| # Print to verify | |
| #print("Updating EMA Embeddings: ", model.ema_embeddings.weight.size()) | |
| #print("Updating Eval Embeddings: ", model.eval_embeddings.weight.size()) | |
| train_params['params']['class_to_taxa'] = train_dataset.class_to_taxa | |
| for _, batch in tqdm(enumerate(train_loader)): | |
| if train_params['params']['use_text_inputs']: | |
| loc_feat, _, class_id, context_feats, _, context_mask, embs = batch | |
| loc_feat = loc_feat.to(eval_params['device']) | |
| class_id = class_id.to(eval_params['device']) | |
| context_feats = context_feats.to(eval_params['device']) | |
| context_mask = context_mask.to(eval_params['device']) | |
| embs = embs.to(eval_params['device']) | |
| # Don't need to do anything with these probs - I am just updating the "eval embeddings" | |
| probs = model.forward( | |
| x=loc_feat, | |
| context_sequence=context_feats, | |
| context_mask=context_mask, | |
| class_ids=class_id, | |
| return_feats=False, | |
| return_class_embeddings=False, | |
| class_of_interest=None, | |
| use_eval_embeddings=True, | |
| text_emb = embs) | |
| else: | |
| loc_feat, _, class_id, context_feats, _, context_mask = batch | |
| loc_feat = loc_feat.to(eval_params['device']) | |
| class_id = class_id.to(eval_params['device']) | |
| context_feats = context_feats.to(eval_params['device']) | |
| context_mask = context_mask.to(eval_params['device']) | |
| # Don't need to do anything with these probs - I am just updating the "eval embeddings" | |
| probs = model.forward( | |
| x=loc_feat, | |
| context_sequence=context_feats, | |
| context_mask=context_mask, | |
| class_ids=class_id, | |
| return_feats=False, | |
| return_class_embeddings=False, | |
| class_of_interest=None, | |
| use_eval_embeddings=True | |
| ) | |
| print('eval embeddings generated!') | |
| elif train_params['params']['model'] == 'VariableInputModel': | |
| train_dataset = datasets.get_train_data(train_params['params']) | |
| if train_dataset.use_text: | |
| if eval_params['text_section'] != '': | |
| train_dataset.select_text_section(eval_params['text_section']) | |
| print(f'Using {eval_params["text_section"]} text for evaluation') | |
| train_loader = torch.utils.data.DataLoader( | |
| train_dataset, | |
| batch_size=train_params['params']['batch_size'], | |
| shuffle=True, | |
| num_workers=8, | |
| collate_fn=getattr(train_dataset, 'collate_fn', None)) | |
| # if len(train_params['params']['class_to_taxa']) != train_dataset.class_to_taxa: | |
| # Create new embedding layers for the expanded classes | |
| num_new_classes = len(train_dataset.class_to_taxa) | |
| embedding_dim = model.ema_embeddings.embedding_dim | |
| new_ema_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| new_eval_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| nn.init.xavier_uniform_(new_ema_embeddings.weight) | |
| nn.init.xavier_uniform_(new_eval_embeddings.weight) | |
| # Convert lists to numpy arrays for indexing | |
| class_to_taxa_np = np.array(train_params['params']['class_to_taxa']) | |
| class_to_taxa_expanded_np = np.array(train_dataset.class_to_taxa) | |
| # Find common taxa and their indices | |
| common_taxa, original_indices, expanded_indices = np.intersect1d( | |
| class_to_taxa_np, class_to_taxa_expanded_np, return_indices=True) | |
| # Update new embeddings for the common taxa | |
| new_ema_embeddings.weight.data[expanded_indices] = model.ema_embeddings.weight.data[original_indices] | |
| new_eval_embeddings.weight.data[expanded_indices] = model.eval_embeddings.weight.data[original_indices] | |
| # Replace old embeddings with new embeddings | |
| model.ema_embeddings = new_ema_embeddings | |
| model.eval_embeddings = new_eval_embeddings | |
| # Print to verify | |
| #print("Updating EMA Embeddings: ", model.ema_embeddings.weight.size()) | |
| #print("Updating Eval Embeddings: ", model.eval_embeddings.weight.size()) | |
| train_params['params']['class_to_taxa'] = train_dataset.class_to_taxa | |
| for _, batch in tqdm(enumerate(train_loader)): | |
| loc_feat, _, class_id, context_feats, _, context_mask, text_emb, image_emb, env_emb = batch | |
| # print('DO I NEED THE BELOW LINES? DO THEY SLOW THINGS DOWN') | |
| # return padded_sequences, padded_locs, class_ids, sequence_mask | |
| loc_feat = loc_feat.to(eval_params['device']) | |
| class_id = class_id.to(eval_params['device']) | |
| context_feats = context_feats.to(eval_params['device']) | |
| context_mask = context_mask.to(eval_params['device']) | |
| text_emb = text_emb.to(eval_params['device']) | |
| image_emb = image_emb.to(eval_params['device']) | |
| if env_emb is not None: | |
| env_emb = env_emb.to(eval_params['device']) | |
| # Don't need to do anything with these probs - I am just updating the "eval embeddings" | |
| probs = model.forward(x=loc_feat, | |
| context_sequence=context_feats, | |
| context_mask=context_mask, | |
| class_ids=class_id, | |
| text_emb=text_emb, | |
| image_emb=image_emb, | |
| env_emb=env_emb, | |
| return_feats=False, | |
| return_class_embeddings=False, | |
| class_of_interest=None, | |
| use_eval_embeddings=True) | |
| print('eval embeddings generated!') | |
| print('\n' + eval_params['eval_type']) | |
| t = time.time() | |
| if eval_params['eval_type'] == 'snt': | |
| eval_params['split'] = 'test' # val, test, all | |
| eval_params['val_frac'] = 0.50 | |
| eval_params['split_seed'] = 7499 | |
| evaluator = EvaluatorSNT(train_params['params'], eval_params) | |
| results = evaluator.run_evaluation(model, enc, extra_input=extra_input) | |
| evaluator.report(results) | |
| elif eval_params['eval_type'] == 'iucn': | |
| evaluator = EvaluatorIUCN(train_params['params'], eval_params) | |
| results = evaluator.run_evaluation(model, enc, extra_input=extra_input) | |
| evaluator.report(results) | |
| elif eval_params['eval_type'] == 'geo_prior': | |
| evaluator = EvaluatorGeoPrior(train_params['params'], eval_params) | |
| results = evaluator.run_evaluation(model, enc, extra_input=extra_input) | |
| evaluator.report(results) | |
| elif eval_params['eval_type'] == 'geo_feature': | |
| evaluator = EvaluatorGeoFeature(train_params['params'], eval_params) | |
| results = evaluator.run_evaluation(model, enc, extra_input=extra_input) | |
| evaluator.report(results) | |
| else: | |
| raise NotImplementedError('Eval type not implemented.') | |
| print(f'evaluation completed in {np.around((time.time()-t)/60, 1)} min') | |
| return results | |
| class EvaluatorGeoPriorLowRank: | |
| def __init__(self, train_params, eval_params): | |
| # store parameters: | |
| self.train_params = train_params | |
| self.eval_params = eval_params | |
| with open('paths.json', 'r') as f: | |
| paths = json.load(f) | |
| # load vision model predictions: | |
| self.data = np.load(os.path.join(paths['geo_prior'], 'geo_prior_model_preds.npz')) | |
| print(self.data['probs'].shape[0], 'total test observations') | |
| # load locations: | |
| meta = pd.read_csv(os.path.join(paths['geo_prior'], 'geo_prior_model_meta.csv')) | |
| self.obs_locs = np.vstack((meta['longitude'].values, meta['latitude'].values)).T.astype(np.float32) | |
| temp = np.array(meta['observed_on'].values, dtype='S10') | |
| temp = temp.view('S1').reshape((temp.size, -1)) | |
| years = temp[:, :4].view('S4').astype(int)[:, 0] | |
| months = temp[:, 5:7].view('S2').astype(int)[:, 0] | |
| days = temp[:, 8:10].view('S2').astype(int)[:, 0] | |
| days_per_month = np.cumsum([0] + [monthrange(2018, mm)[1] for mm in range(1, 12)]) | |
| dates = days_per_month[months - 1] + days - 1 | |
| self.dates = np.round((dates) / 365.0, 4).astype(np.float32) | |
| # taxonomic mapping: | |
| self.taxon_map = self.find_mapping_between_models(self.data['model_to_taxa'], self.train_params['class_to_taxa']) | |
| print(self.taxon_map.shape[0], 'out of', len(self.data['model_to_taxa']), 'taxa in both vision and geo models') | |
| def find_mapping_between_models(self, vision_taxa, geo_taxa): | |
| # this will output an array of size N_overlap X 2 | |
| # the first column will be the indices of the vision model, and the second is their | |
| # corresponding index in the geo model | |
| taxon_map = np.ones((vision_taxa.shape[0], 2), dtype=np.int32)*-1 | |
| taxon_map[:, 0] = np.arange(vision_taxa.shape[0]) | |
| geo_taxa_arr = np.array(geo_taxa) | |
| for tt_id, tt in enumerate(vision_taxa): | |
| ind = np.where(geo_taxa_arr==tt)[0] | |
| if len(ind) > 0: | |
| taxon_map[tt_id, 1] = ind[0] | |
| inds = np.where(taxon_map[:, 1]>-1)[0] | |
| taxon_map = taxon_map[inds, :] | |
| return taxon_map | |
| def convert_to_inat_vision_order(self, geo_pred_ip, vision_top_k_prob, vision_top_k_inds, vision_taxa, taxon_map): | |
| # this is slow as we turn the sparse input back into the same size as the dense one | |
| vision_pred = np.zeros((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32) | |
| geo_pred = np.ones((geo_pred_ip.shape[0], len(vision_taxa)), dtype=np.float32) | |
| vision_pred[np.arange(vision_pred.shape[0])[..., np.newaxis], vision_top_k_inds] = vision_top_k_prob | |
| geo_pred[:, taxon_map[:, 0]] = geo_pred_ip[:, taxon_map[:, 1]] | |
| return geo_pred, vision_pred | |
| def run_evaluation(self, model): | |
| results = {} | |
| # loop over in batches | |
| batch_start = np.hstack((np.arange(0, self.data['probs'].shape[0], self.eval_params['batch_size']), self.data['probs'].shape[0])) | |
| correct_pred = np.zeros(self.data['probs'].shape[0]) | |
| from tqdm import tqdm | |
| for bb_id, bb in tqdm(enumerate(range(len(batch_start)-1))): | |
| batch_inds = np.arange(batch_start[bb], batch_start[bb+1]) | |
| vision_probs = self.data['probs'][batch_inds, :] | |
| vision_inds = self.data['inds'][batch_inds, :] | |
| gt = self.data['labels'][batch_inds] | |
| dates = torch.from_numpy(self.dates[batch_inds]) | |
| obs_locs_batch = torch.from_numpy(self.obs_locs[batch_inds, :]).to(self.eval_params['device']) | |
| with torch.no_grad(): | |
| geo_pdf = torch.log(model.sample(obs_locs_batch)).T | |
| for bias in range(11+5, 12+5): | |
| geo_pred, vision_pred = self.convert_to_inat_vision_order(geo_pdf+bias, vision_probs, vision_inds, | |
| self.data['model_to_taxa'], self.taxon_map) | |
| geo_pred = softmax(torch.from_numpy(geo_pred), dim=1).numpy() | |
| #print(bias, (np.argmax(vision_pred*geo_pred2, 1) == gt).mean().item()) | |
| comb_pred = np.argmax(vision_pred*geo_pred, 1) | |
| comb_pred = (comb_pred==gt) | |
| correct_pred[batch_inds] = comb_pred | |
| accuracy_by_taxa = np.zeros(len(self.data['model_to_taxa'])) | |
| for tt_id, tt in enumerate(self.data['model_to_taxa']): | |
| inds = np.where(self.data['labels'] == tt)[0] | |
| accuracy_by_taxa[tt_id] = float((correct_pred[inds].mean())) | |
| results['vision_only_top_1'] = float((self.data['inds'][:, -1] == self.data['labels']).mean()) | |
| results['vision_geo_top_1'] = float(correct_pred.mean()) | |
| return results | |
| def report(self, results): | |
| print('Overall accuracy vision only model', round(results['vision_only_top_1'], 3)) | |
| print('Overall accuracy of geo model ', round(results['vision_geo_top_1'], 3)) | |
| print('Gain ', round(results['vision_geo_top_1'] - results['vision_only_top_1'], 3)) | |
| # MINE MINE MINE - these are just to help with low shot plotting. Can probably be elsewhere. | |
| def generate_eval_embeddings(overrides, taxa_of_interest, num_context, train_overrides=None): | |
| eval_params = setup.get_default_params_eval(overrides) | |
| # set up model: | |
| eval_params['model_path'] = os.path.join(eval_params['exp_base'], eval_params['experiment_name'], eval_params['ckp_name']) | |
| eval_params['device'] = 'cpu' | |
| train_params = torch.load(eval_params['model_path'], map_location='cpu') | |
| train_params['params']['device'] = 'cpu' | |
| default_params = setup.get_default_params_train() | |
| for key in default_params: | |
| if key not in train_params['params']: | |
| train_params['params'][key] = default_params[key] | |
| # create input encoder: | |
| if train_params['params']['input_enc'] in ['env', 'sin_cos_env']: | |
| raster = datasets.load_env().to(eval_params['device']) | |
| else: | |
| raster = None | |
| enc = utils.CoordEncoder(train_params['params']['input_enc'], raster=raster, input_dim=train_params['params']['input_dim']) | |
| if train_params['params']['input_time']: | |
| time_enc = utils.TimeEncoder(input_enc='conical') if train_params['params']['input_time'] else None | |
| extra_input = torch.cat([time_enc.encode(torch.tensor([[0.0, 1.0]]))], dim=1).to(eval_params['device']) | |
| else: | |
| extra_input = None | |
| if train_overrides != None: | |
| for key, value in train_overrides.items(): | |
| #print(f'updating train param {key}') | |
| train_params['params'][key] = value | |
| train_dataset = datasets.get_train_data(train_params['params']) | |
| train_loader = torch.utils.data.DataLoader( | |
| train_dataset, | |
| batch_size=train_params['params']['batch_size'], | |
| shuffle=True, | |
| num_workers=8, | |
| collate_fn=getattr(train_dataset, 'collate_fn', None)) | |
| model = models.get_model(train_params['params'], inference_only=True) | |
| # model.load_state_dict(train_params['state_dict'], strict=True) | |
| model.load_state_dict(train_params['state_dict'], strict=False) | |
| model = model.to(eval_params['device']) | |
| model.eval() | |
| # Create new embedding layers for the expanded classes | |
| num_new_classes = len(train_dataset.class_to_taxa) | |
| embedding_dim = model.ema_embeddings.embedding_dim | |
| new_ema_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| new_eval_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| nn.init.xavier_uniform_(new_ema_embeddings.weight) | |
| nn.init.xavier_uniform_(new_eval_embeddings.weight) | |
| # Convert lists to numpy arrays for indexing | |
| class_to_taxa_np = np.array(train_params['params']['class_to_taxa']) | |
| class_to_taxa_expanded_np = np.array(train_dataset.class_to_taxa) | |
| # Find common taxa and their indices | |
| common_taxa, original_indices, expanded_indices = np.intersect1d( | |
| class_to_taxa_np, class_to_taxa_expanded_np, return_indices=True) | |
| # Update new embeddings for the common taxa | |
| new_ema_embeddings.weight.data[expanded_indices] = model.ema_embeddings.weight.data[original_indices] | |
| new_eval_embeddings.weight.data[expanded_indices] = model.eval_embeddings.weight.data[original_indices] | |
| # Replace old embeddings with new embeddings | |
| model.ema_embeddings = new_ema_embeddings | |
| model.eval_embeddings = new_eval_embeddings | |
| # Print to verify | |
| #print("Updated EMA Embeddings: ", model.ema_embeddings.weight.size()) | |
| #print("Updated Eval Embeddings: ", model.eval_embeddings.weight.size()) | |
| train_params['params']['class_to_taxa'] = train_dataset.class_to_taxa | |
| class_of_interest = train_dataset.class_to_taxa.index(taxa_of_interest) | |
| # Find the index of class_of_interest in the labels tensor | |
| loc_index_of_interest = (train_dataset.labels == class_of_interest).nonzero(as_tuple=True)[0].item() | |
| # loc_index_of_interest = train_dataset.labels.index(class_of_interest) | |
| loc_of_interest = train_dataset.loc_feats[loc_index_of_interest] | |
| all_class_context_feats = train_dataset.per_class_loc_feats[class_of_interest] | |
| all_class_context_locs = train_dataset.per_class_locs[class_of_interest] | |
| context_feats_of_interest = all_class_context_feats[:num_context,:] | |
| context_locs_of_interest = all_class_context_locs[:num_context,:] | |
| # context_mask = context_feats_of_interest != -10 | |
| # context_mask = None | |
| # context_mask = (context_locs_of_interest == -10).all(dim=-1).to(eval_params['device']) | |
| context_mask = (context_locs_of_interest == -10).all(dim=-1).to(eval_params['device']).unsqueeze(0) | |
| probs = model.forward( | |
| x=loc_of_interest.to(train_params['params']['device']), | |
| context_sequence=context_feats_of_interest.to(train_params['params']['device']), | |
| context_mask=context_mask, | |
| class_ids=class_of_interest, | |
| return_feats=False, | |
| return_class_embeddings=False, | |
| class_of_interest=None, | |
| use_eval_embeddings=True | |
| ) | |
| #print(f'eval embedding generated for class {class_of_interest}, taxa {taxa_of_interest}') | |
| return model, context_locs_of_interest, train_params, class_of_interest | |
| def generate_eval_embedding_from_given_points(context_points, overrides, taxa_of_interest, train_overrides=None, text_emb=None): | |
| eval_params = setup.get_default_params_eval(overrides) | |
| # set up model: | |
| eval_params['model_path'] = os.path.join(eval_params['exp_base'], eval_params['experiment_name'], eval_params['ckp_name']) | |
| train_params = torch.load(eval_params['model_path'], map_location='cpu') | |
| default_params = setup.get_default_params_train() | |
| for key in default_params: | |
| if key not in train_params['params']: | |
| train_params['params'][key] = default_params[key] | |
| # create input encoder: | |
| if train_params['params']['input_enc'] in ['env', 'sin_cos_env']: | |
| raster = datasets.load_env().to(eval_params['device']) | |
| else: | |
| raster = None | |
| enc = utils.CoordEncoder(train_params['params']['input_enc'], raster=raster, input_dim=train_params['params']['input_dim']) | |
| if train_params['params']['input_time']: | |
| time_enc = utils.TimeEncoder(input_enc='conical') if train_params['params']['input_time'] else None | |
| extra_input = torch.cat([time_enc.encode(torch.tensor([[0.0, 1.0]]))], dim=1).to(eval_params['device']) | |
| else: | |
| extra_input = None | |
| if train_overrides != None: | |
| for key, value in train_overrides.items(): | |
| #print(f'updating train param {key}') | |
| train_params['params'][key] = value | |
| # create context point encoder | |
| transformer_input_enc = train_params['params']['transformer_input_enc'] | |
| if transformer_input_enc in ['env', 'sin_cos_env']: | |
| transformer_raster = datasets.load_env().to(eval_params['device']) | |
| else: | |
| transformer_raster = None | |
| token_dim = train_params['params']['species_dim'] | |
| if transformer_input_enc == 'sinr': | |
| transformer_enc = enc | |
| else: | |
| transformer_enc = utils.CoordEncoder(transformer_input_enc, transformer_raster, input_dim=token_dim) | |
| # transformer_enc = utils.CoordEncoder(transformer_input_enc, transformer_raster, input_dim=token_dim) | |
| # load model | |
| model = models.get_model(train_params['params'], inference_only=True) | |
| # model.load_state_dict(train_params['state_dict'], strict=True) | |
| model.load_state_dict(train_params['state_dict'], strict=False) | |
| model = model.to(eval_params['device']) | |
| model.eval() | |
| # # Create new embedding layers for the expanded classes | |
| # num_new_classes = len(train_params['params']['class_to_taxa']) | |
| embedding_dim = model.ema_embeddings.embedding_dim | |
| # new_ema_embeddings = nn.Embedding(num_embeddings=num_new_classes, embedding_dim=embedding_dim).to(eval_params["device"]) | |
| new_eval_embeddings = nn.Embedding(num_embeddings=model.eval_embeddings.weight.size()[0], embedding_dim=embedding_dim).to(eval_params["device"]) | |
| # Update new embeddings for the common taxa | |
| new_eval_embeddings.weight.data = model.eval_embeddings.weight.data | |
| # Replace old embeddings with new embeddings | |
| model.eval_embeddings = new_eval_embeddings | |
| # Print to verify | |
| #print("Updated EMA Embeddings: ", model.ema_embeddings.weight.size()) | |
| #print("Updated Eval Embeddings: ", model.eval_embeddings.weight.size()) | |
| class_of_interest = 0 | |
| just_loc = torch.from_numpy(np.array([[0.0,0.0]]).astype(np.float32)) | |
| loc_of_interest = enc.encode(just_loc, normalize=False) | |
| context_points = torch.from_numpy(np.array(context_points).astype(np.float32)) | |
| all_class_context_feats = transformer_enc.encode(context_points, normalize=False) | |
| all_class_context_locs = context_points | |
| context_feats_of_interest = all_class_context_feats | |
| context_locs_of_interest = all_class_context_locs | |
| # context_mask = context_feats_of_interest[:,0] != -10 | |
| # context_mask = None | |
| context_mask = torch.from_numpy(np.full((1, context_feats_of_interest.shape[0]), False)) | |
| # probs = model.forward( | |
| # x=loc_of_interest.to(train_params['params']['device']), | |
| # context_sequence=context_feats_of_interest.to(train_params['params']['device']), | |
| # context_mask=context_mask, | |
| # class_ids=class_of_interest, | |
| # return_feats=False, | |
| # return_class_embeddings=False, | |
| # class_of_interest=None, | |
| # use_eval_embeddings=True | |
| # ) | |
| probs = model.forward( | |
| x=loc_of_interest.to(eval_params['device']), | |
| context_sequence=context_feats_of_interest.to(eval_params['device']), | |
| context_mask=context_mask, | |
| class_ids=class_of_interest, | |
| return_feats=False, | |
| return_class_embeddings=False, | |
| class_of_interest=None, | |
| use_eval_embeddings=True, | |
| text_emb=text_emb | |
| ) | |
| #print(f'eval embedding generated for class {class_of_interest}, from hand selected context points') | |
| return model, context_locs_of_interest, train_params, class_of_interest |