def se_core_spatialXfeature_readout(
dataloaders,
seed,
hidden_channels=32,
input_kern=13, # core args
hidden_kern=3,
layers=3,
gamma_input=15.5,
skip=0,
final_nonlinearity=True,
momentum=0.9,
pad_input=False,
batch_norm=True,
hidden_dilation=1,
laplace_padding=None,
input_regularizer='LaplaceL2norm',
init_noise=1e-3,
readout_bias=True, # readout args,
gamma_readout=4,
normalize=False,
elu_offset=0,
stack=None,
se_reduction=32,
n_se_blocks=1,
depth_separable=False,
linear=False):
"""
Model class of a stacked2dCore (from neuralpredictors) and a spatialXfeature (factorized) readout
Args:
Returns: An initialized model which consists of model.core and model.readout
"""
if "train" in dataloaders.keys():
dataloaders = dataloaders["train"]
# Obtain the named tuple fields from the first entry of the first dataloader in the dictionary
in_name, out_name = next(iter(list(dataloaders.values())[0]))._fields
session_shape_dict = get_dims_for_loader_dict(dataloaders)
n_neurons_dict = {k: v[out_name][1] for k, v in session_shape_dict.items()}
in_shapes_dict = {k: v[in_name] for k, v in session_shape_dict.items()}
input_channels = [v[in_name][1] for v in session_shape_dict.values()]
class Encoder(nn.Module):
def __init__(self, core, readout, elu_offset):
super().__init__()
self.core = core
self.readout = readout
self.offset = elu_offset
def forward(self, x, data_key=None, **kwargs):
x = self.core(x)
x = self.readout(
x,
data_key=data_key,
)
return F.elu(x + self.offset) + 1
def regularizer(self, data_key):
return self.core.regularizer() + self.readout.regularizer(
data_key=data_key)
set_random_seed(seed)
# get a stacked2D core from neuralpredictors
core = SE2dCore(input_channels=input_channels[0],
hidden_channels=hidden_channels,
input_kern=input_kern,
hidden_kern=hidden_kern,
layers=layers,
gamma_input=gamma_input,
skip=skip,
final_nonlinearity=final_nonlinearity,
bias=False,
momentum=momentum,
pad_input=pad_input,
batch_norm=batch_norm,
hidden_dilation=hidden_dilation,
laplace_padding=laplace_padding,
input_regularizer=input_regularizer,
stack=stack,
se_reduction=se_reduction,
n_se_blocks=n_se_blocks,
depth_separable=depth_separable,
linear=linear)
readout = MultipleSpatialXFeatureLinear(core,
in_shape_dict=in_shapes_dict,
n_neurons_dict=n_neurons_dict,
init_noise=init_noise,
bias=readout_bias,
gamma_readout=gamma_readout,
normalize=normalize)
# initializing readout bias to mean response
if readout_bias:
for key, value in dataloaders.items():
_, targets = next(iter(value))
readout[key].bias.data = targets.mean(0)
model = Encoder(core, readout, elu_offset)
return model
def neurips_loader(path,
batch_size,
seed=None,
areas=None,
layers=None,
tier=None,
neuron_ids=None,
get_key=False,
cuda=True,
normalize=True,
exclude=None,
trial_ids=None,
**kwargs):
dat = FileTreeDataset(path, 'images', 'responses')
# The permutation MUST be added first and the conditions below MUST NOT be based on the original order
# specify condition(s) for sampling neurons. If you want to sample specific neurons define conditions that would effect idx
conds = np.ones(len(dat.neurons.area), dtype=bool)
if areas is not None:
conds &= (np.isin(dat.neurons.area, areas))
if layers is not None:
conds &= (np.isin(dat.neurons.layer, layers))
if neuron_ids is not None:
conds &= (np.isin(dat.neurons.unit_ids, neuron_ids))
idx = np.where(conds)[0]
more_transforms = [Subsample(idx), ToTensor(cuda)]
if normalize:
more_transforms.insert(0, NeuroNormalizer(dat, exclude=exclude))
dat.transforms.extend(more_transforms)
# subsample images
dataloaders = {}
keys = [tier] if tier else ['train', 'validation', 'test']
for tier in keys:
if seed is not None:
set_random_seed(seed)
# torch.manual_seed(img_seed)
# sample images
conds = np.ones_like(dat.trial_info.tiers).astype(bool)
conds &= (dat.trial_info.tiers == tier)
if trial_ids is not None and tier in trial_ids:
print(f'Subsampling {tier} set to {len(trial_ids[tier])} trials',
flush=True)
conds &= (np.isin(dat.trial_info.trial_idx, trial_ids[tier]))
subset_idx = np.where(conds)[0]
if tier == 'train':
sampler = SubsetRandomSampler(subset_idx)
dataloaders[tier] = DataLoader(dat,
sampler=sampler,
batch_size=batch_size)
elif tier == 'validation':
sampler = SubsetSequentialSampler(subset_idx)
dataloaders[tier] = DataLoader(dat,
sampler=sampler,
batch_size=batch_size)
else:
sampler = RepeatsBatchSampler(
dataloaders['train'].dataset.trial_info.frame_image_id,
subset_index=subset_idx)
dataloaders[tier] = DataLoader(dat,
batch_sampler=sampler,
sampler=None)
# create the data_key for a specific data path
data_key = path.split('static')[-1].split('.')[0].replace('preproc', '')
return (data_key, dataloaders) if get_key else dataloaders
def nnvision_trainer(
model,
dataloaders,
seed,
avg_loss=False,
scale_loss=True, # trainer args
loss_function='PoissonLoss',
stop_function='get_correlations',
loss_accum_batch_n=None,
device='cuda',
verbose=True,
interval=1,
patience=5,
epoch=0,
lr_init=0.005, # early stopping args
max_iter=100,
maximize=True,
tolerance=1e-6,
restore_best=True,
lr_decay_steps=3,
lr_decay_factor=0.3,
min_lr=0.0001, # lr scheduler args
cb=None,
track_training=False,
return_test_score=False,
**kwargs):
"""
Args:
model:
dataloaders:
seed:
avg_loss:
scale_loss:
loss_function:
stop_function:
loss_accum_batch_n:
device:
verbose:
interval:
patience:
epoch:
lr_init:
max_iter:
maximize:
tolerance:
restore_best:
lr_decay_steps:
lr_decay_factor:
min_lr:
cb:
track_training:
**kwargs:
Returns:
"""
def full_objective(model, dataloader, data_key, *args):
"""
Args:
model:
dataloader:
data_key:
*args:
Returns:
"""
loss_scale = np.sqrt(
len(dataloader[data_key].dataset) /
args[0].shape[0]) if scale_loss else 1.0
return loss_scale * criterion(model(args[0].to(device), data_key), args[1].to(device)) \
+ model.regularizer(data_key)
##### Model training ####################################################################################################
model.to(device)
set_random_seed(seed)
model.train()
criterion = getattr(mlmeasures, loss_function)(avg=avg_loss)
stop_closure = partial(getattr(measures, stop_function),
dataloaders=dataloaders["validation"],
device=device,
per_neuron=False,
avg=True)
n_iterations = len(LongCycler(dataloaders["train"]))
optimizer = torch.optim.Adam(model.parameters(), lr=lr_init)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='max' if maximize else 'min',
factor=lr_decay_factor,
patience=patience,
threshold=tolerance,
min_lr=min_lr,
verbose=verbose,
threshold_mode='abs')
# set the number of iterations over which you would like to accummulate gradients
optim_step_count = len(dataloaders["train"].keys(
)) if loss_accum_batch_n is None else loss_accum_batch_n
if track_training:
tracker_dict = dict(correlation=partial(get_correlations(),
model,
dataloaders["validation"],
device=device,
per_neuron=False),
poisson_loss=partial(get_poisson_loss(),
model,
dataloaders["validation"],
device=device,
per_neuron=False,
avg=False))
if hasattr(model, 'tracked_values'):
tracker_dict.update(model.tracked_values)
tracker = MultipleObjectiveTracker(**tracker_dict)
else:
tracker = None
# train over epochs
for epoch, val_obj in early_stopping(model,
stop_closure,
interval=interval,
patience=patience,
start=epoch,
max_iter=max_iter,
maximize=maximize,
tolerance=tolerance,
restore_best=restore_best,
tracker=tracker,
scheduler=scheduler,
lr_decay_steps=lr_decay_steps):
# print the quantities from tracker
if verbose and tracker is not None:
print("=======================================")
for key in tracker.log.keys():
print(key, tracker.log[key][-1], flush=True)
# executes callback function if passed in keyword args
if cb is not None:
cb()
# train over batches
optimizer.zero_grad()
for batch_no, (data_key, data) in tqdm(enumerate(
LongCycler(dataloaders["train"])),
total=n_iterations,
desc="Epoch {}".format(epoch)):
loss = full_objective(model, dataloaders["train"], data_key, *data)
loss.backward()
if (batch_no + 1) % optim_step_count == 0:
optimizer.step()
optimizer.zero_grad()
##### Model evaluation ####################################################################################################
model.eval()
tracker.finalize() if track_training else None
# Compute avg validation and test correlation
validation_correlation = get_correlations(model,
dataloaders["validation"],
device=device,
as_dict=False,
per_neuron=False)
test_correlation = get_correlations(model,
dataloaders["test"],
device=device,
as_dict=False,
per_neuron=False)
# return the whole tracker output as a dict
output = {k: v for k, v in tracker.log.items()} if track_training else {}
output["validation_corr"] = validation_correlation
score = np.mean(test_correlation) if return_test_score else np.mean(
validation_correlation)
return score, output, model.state_dict()
def static_shared_loaders(
paths,
batch_size,
seed=None,
areas=None,
layers=None,
tier=None,
multi_match_ids=None,
multi_match_n=None,
exclude_multi_match_n=0,
multi_match_base_seed=None,
image_ids=None,
image_n=None,
image_base_seed=None,
cuda=True,
normalize=True,
include_behavior=False,
exclude="images",
select_input_channel=None,
return_test_sampler=False,
oracle_condition=None,
):
"""
Returns a dictionary of dataloaders (i.e., trainloaders, valloaders, and testloaders) for >= 1 dataset(s).
The datasets must have matched neurons. Only the file tree format is supported.
Args:
paths (list): list of paths for the datasets
batch_size (int): batch size.
seed (int): seed. Not really needed because there are neuron and image seed. But nnFabrik requires it.
areas (list, optional): the visual area.
layers (list, optional): the layer from visual area.
tier (str, optional): tier is a placeholder to specify which set of images to pick for train, val, and test loader.
multi_match_ids (list, optional): List of multi_match_ids according to which the respective neuron_ids are drawn for each dataset in paths
multi_match_n (int, optional): number of neurons to select randomly. Can not be set together with multi_match_ids
exclude_multi_match_n (int): the first <exclude_multi_match_n> matched neurons will be excluded (given a multi_match_base_seed),
then <multi_match_n> matched neurons will be drawn from the remaining neurons.
multi_match_base_seed (float, optional): base seed for neuron selection. Get's multiplied by multi_match_n to obtain final seed
image_ids (list, optional): List of lists of image_ids. Make sure the order is the same as in paths
image_n (int, optional): number of images to select randomly. Can not be set together with image_ids
image_base_seed (float, optional): base seed for image selection. Get's multiplied by image_n to obtain final seed
cuda (bool, optional): whether to place the data on gpu or not.
normalize (bool, optional): whether to normalize the data (see also exclude)
exclude (str, optional): data to exclude from data-normalization. Only relevant if normalize=True. Defaults to 'images'
include_behavior (bool, optional): whether to include behavioral data
select_input_channel (int, optional): Only for color images. Select a color channel
return_test_sampler (bool, optional): whether to return only the test loader with repeat-batches
oracle_condition (list, optional): Only relevant if return_test_sampler=True. Class indices for the sampler
Returns:
dict: dictionary of dictionaries where the first level keys are 'train', 'validation', and 'test', and second level keys are data_keys.
"""
set_random_seed(seed)
assert (
len(paths) != 1
), "Only one dataset was specified in 'paths'. When using the 'static_shared_loaders', more than one dataset has to be passed."
assert any(
[
multi_match_ids is None,
all([
multi_match_n is None, multi_match_base_seed is None,
exclude_multi_match_n == 0
]),
]
), "multi_match_ids can not be set at the same time with any other multi_match selection criteria"
assert any([
exclude_multi_match_n == 0, multi_match_base_seed is not None
]), "multi_match_base_seed must be set when exclude_multi_match_n is not 0"
# Collect overlapping multi matches
multi_unit_ids, per_data_set_ids, given_neuron_ids = [], [], []
match_set = None
for path in paths:
data_key, dataloaders = static_loader(path=path,
batch_size=100,
get_key=True)
dat = dataloaders["train"].dataset
multi_unit_ids.append(dat.neurons.multi_match_id)
per_data_set_ids.append(dat.neurons.unit_ids)
if match_set is None:
match_set = set(multi_unit_ids[-1])
else:
match_set &= set(multi_unit_ids[-1])
if multi_match_ids is not None:
assert set(multi_match_ids).issubset(
dat.neurons.multi_match_id
), "Dataset {} does not contain all multi_match_ids".format(path)
neuron_idx = [
np.where(dat.neurons.multi_match_id == multi_match_id)[0][0]
for multi_match_id in multi_match_ids
]
given_neuron_ids.append(dat.neurons.unit_ids[neuron_idx])
match_set -= {-1} # remove the unmatched neurons
match_set = np.array(list(match_set))
# get unit_ids of matched neurons
if multi_match_ids is not None:
neuron_ids = given_neuron_ids
elif multi_match_n is not None:
random_state = np.random.get_state()
if multi_match_base_seed is not None:
np.random.seed(
multi_match_base_seed * multi_match_n
) # avoid nesting by making seed dependent on number of neurons
assert (
len(match_set) >= exclude_multi_match_n + multi_match_n
), "After excluding {} neurons, there are not {} matched neurons left".format(
exclude_multi_match_n, multi_match_n)
match_subset = np.random.choice(match_set,
size=exclude_multi_match_n +
multi_match_n,
replace=False)[exclude_multi_match_n:]
neuron_ids = [
pdsi[np.isin(munit_ids, match_subset)]
for munit_ids, pdsi in zip(multi_unit_ids, per_data_set_ids)
]
np.random.set_state(random_state)
else:
neuron_ids = [
pdsi[np.isin(munit_ids, match_set)]
for munit_ids, pdsi in zip(multi_unit_ids, per_data_set_ids)
]
# generate single dataloaders
dls = OrderedDict({})
keys = [tier] if tier else ["train", "validation", "test"]
for key in keys:
dls[key] = OrderedDict({})
image_ids = [image_ids] if image_ids is None else image_ids
for path, neuron_id, image_id in zip_longest(paths,
neuron_ids,
image_ids,
fillvalue=None):
data_key, loaders = static_loader(
path,
batch_size,
areas=areas,
layers=layers,
cuda=cuda,
tier=tier,
get_key=True,
neuron_ids=neuron_id,
neuron_n=None,
neuron_base_seed=None,
image_ids=image_id,
image_n=image_n,
image_base_seed=image_base_seed,
normalize=normalize,
include_behavior=include_behavior,
exclude=exclude,
select_input_channel=select_input_channel,
file_tree=True,
return_test_sampler=return_test_sampler,
oracle_condition=oracle_condition,
)
for k in dls:
dls[k][data_key] = loaders[k]
return dls
def standard_trainer(model,
dataloaders,
seed,
avg_loss=False,
scale_loss=True,
loss_function="PoissonLoss",
stop_function="get_correlations",
loss_accum_batch_n=None,
device="cuda",
verbose=True,
interval=1,
patience=5,
epoch=0,
lr_init=0.005,
max_iter=200,
maximize=True,
tolerance=1e-6,
restore_best=True,
lr_decay_steps=3,
lr_decay_factor=0.3,
min_lr=0.0001,
cb=None,
track_training=False,
return_test_score=False,
detach_core=False,
**kwargs):
"""
Args:
model: model to be trained
dataloaders: dataloaders containing the data to train the model with
seed: random seed
avg_loss: whether to average (or sum) the loss over a batch
scale_loss: whether to scale the loss according to the size of the dataset
loss_function: loss function to use
stop_function: the function (metric) that is used to determine the end of the training in early stopping
loss_accum_batch_n: number of batches to accumulate the loss over
device: device to run the training on
verbose: whether to print out a message for each optimizer step
interval: interval at which objective is evaluated to consider early stopping
patience: number of times the objective is allowed to not become better before the iterator terminates
epoch: starting epoch
lr_init: initial learning rate
max_iter: maximum number of training iterations
maximize: whether to maximize or minimize the objective function
tolerance: tolerance for early stopping
restore_best: whether to restore the model to the best state after early stopping
lr_decay_steps: how many times to decay the learning rate after no improvement
lr_decay_factor: factor to decay the learning rate with
min_lr: minimum learning rate
cb: whether to execute callback function
track_training: whether to track and print out the training progress
**kwargs:
Returns:
"""
def full_objective(model, dataloader, data_key, *args, detach_core):
loss_scale = np.sqrt(
len(dataloader[data_key].dataset) /
args[0].shape[0]) if scale_loss else 1.0
regularizers = int(not detach_core) * model.core.regularizer(
) + model.readout.regularizer(data_key)
return (loss_scale * criterion(
model(args[0].to(device), data_key, detach_core=detach_core),
args[1].to(device)) + regularizers)
##### Model training ####################################################################################################
model.to(device)
set_random_seed(seed)
model.train()
criterion = getattr(mlmeasures, loss_function)(avg=avg_loss)
stop_closure = partial(
getattr(measures, stop_function),
dataloaders=dataloaders["validation"],
device=device,
per_neuron=False,
avg=True,
)
n_iterations = len(LongCycler(dataloaders["train"]))
optimizer = torch.optim.Adam(model.parameters(), lr=lr_init)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode="max" if maximize else "min",
factor=lr_decay_factor,
patience=patience,
threshold=tolerance,
min_lr=min_lr,
verbose=verbose,
threshold_mode="abs",
)
# set the number of iterations over which you would like to accummulate gradients
optim_step_count = len(dataloaders["train"].keys(
)) if loss_accum_batch_n is None else loss_accum_batch_n
if track_training:
tracker_dict = dict(
correlation=partial(get_correlations,
model,
dataloaders["validation"],
device=device,
per_neuron=False),
poisson_loss=partial(get_poisson_loss,
model,
dataloaders["validation"],
device=device,
per_neuron=False,
avg=False),
)
if hasattr(model, "tracked_values"):
tracker_dict.update(model.tracked_values)
tracker = MultipleObjectiveTracker(**tracker_dict)
else:
tracker = None
# train over epochs
for epoch, val_obj in early_stopping(
model,
stop_closure,
interval=interval,
patience=patience,
start=epoch,
max_iter=max_iter,
maximize=maximize,
tolerance=tolerance,
restore_best=restore_best,
tracker=tracker,
scheduler=scheduler,
lr_decay_steps=lr_decay_steps,
):
# print the quantities from tracker
if verbose and tracker is not None:
print("=======================================")
for key in tracker.log.keys():
print(key, tracker.log[key][-1], flush=True)
# executes callback function if passed in keyword args
if cb is not None:
cb()
# train over batches
optimizer.zero_grad()
for batch_no, (data_key, data) in tqdm(enumerate(
LongCycler(dataloaders["train"])),
total=n_iterations,
desc="Epoch {}".format(epoch)):
loss = full_objective(model,
dataloaders["train"],
data_key,
*data,
detach_core=detach_core)
loss.backward()
if (batch_no + 1) % optim_step_count == 0:
optimizer.step()
optimizer.zero_grad()
##### Model evaluation ####################################################################################################
model.eval()
tracker.finalize() if track_training else None
# Compute avg validation and test correlation
validation_correlation = get_correlations(model,
dataloaders["validation"],
device=device,
as_dict=False,
per_neuron=False)
test_correlation = get_correlations(model,
dataloaders["test"],
device=device,
as_dict=False,
per_neuron=False)
# return the whole tracker output as a dict
output = {k: v for k, v in tracker.log.items()} if track_training else {}
output["validation_corr"] = validation_correlation
score = np.mean(test_correlation) if return_test_score else np.mean(
validation_correlation)
return score, output, model.state_dict()
def static_loaders(
paths,
batch_size,
seed=None,
areas=None,
layers=None,
tier=None,
neuron_ids=None,
neuron_n=None,
exclude_neuron_n=0,
neuron_base_seed=None,
image_ids=None,
image_n=None,
image_base_seed=None,
cuda=True,
normalize=True,
include_behavior=False,
exclude="images",
select_input_channel=None,
file_tree=True,
return_test_sampler=False,
oracle_condition=None,
):
"""
Returns a dictionary of dataloaders (i.e., trainloaders, valloaders, and testloaders) for >= 1 dataset(s).
Args:
paths (list): list of paths for the datasets
batch_size (int): batch size.
seed (int): seed. Not really needed because there are neuron and image seed. But nnFabrik requires it.
areas (list, optional): the visual area.
layers (list, optional): the layer from visual area.
tier (str, optional): tier is a placeholder to specify which set of images to pick for train, val, and test loader.
neuron_ids (list, optional): List of lists of neuron_ids. Make sure the order is the same as in paths
neuron_n (int, optional): number of neurons to select randomly. Can not be set together with neuron_ids
exclude_neuron_n (int): the first <exclude_neuron_n> neurons will be excluded (given a neuron_base_seed),
then <neuron_n> neurons will be drawn from the remaining neurons.
neuron_base_seed (float, optional): base seed for neuron selection. Get's multiplied by neuron_n to obtain final seed
image_ids (list, optional): List of lists of image_ids. Make sure the order is the same as in paths
image_n (int, optional): number of images to select randomly. Can not be set together with image_ids
image_base_seed (float, optional): base seed for image selection. Get's multiplied by image_n to obtain final seed
cuda (bool, optional): whether to place the data on gpu or not.
normalize (bool, optional): whether to normalize the data (see also exclude)
exclude (str, optional): data to exclude from data-normalization. Only relevant if normalize=True. Defaults to 'images'
include_behavior (bool, optional): whether to include behavioral data
select_input_channel (int, optional): Only for color images. Select a color channel
file_tree (bool, optional): whether to use the file tree dataset format. If False, equivalent to the HDF5 format
return_test_sampler (bool, optional): whether to return only the test loader with repeat-batches
oracle_condition (list, optional): Only relevant if return_test_sampler=True. Class indices for the sampler
Returns:
dict: dictionary of dictionaries where the first level keys are 'train', 'validation', and 'test', and second level keys are data_keys.
"""
set_random_seed(seed)
dls = OrderedDict({})
keys = [tier] if tier else ["train", "validation", "test"]
for key in keys:
dls[key] = OrderedDict({})
neuron_ids = [neuron_ids] if neuron_ids is None else neuron_ids
image_ids = [image_ids] if image_ids is None else image_ids
for path, neuron_id, image_id in zip_longest(paths,
neuron_ids,
image_ids,
fillvalue=None):
data_key, loaders = static_loader(
path,
batch_size,
areas=areas,
layers=layers,
cuda=cuda,
tier=tier,
get_key=True,
neuron_ids=neuron_id,
neuron_n=neuron_n,
exclude_neuron_n=exclude_neuron_n,
neuron_base_seed=neuron_base_seed,
image_ids=image_id,
image_n=image_n,
image_base_seed=image_base_seed,
normalize=normalize,
include_behavior=include_behavior,
exclude=exclude,
select_input_channel=select_input_channel,
file_tree=file_tree,
return_test_sampler=return_test_sampler,
oracle_condition=oracle_condition,
)
for k in dls:
dls[k][data_key] = loaders[k]
return dls
def taskdriven_fullSXF(
dataloaders,
seed,
elu_offset=0,
data_info=None,
# core args
tl_model_name="vgg16",
layers=4,
pretrained=True,
final_batchnorm=True,
final_nonlinearity=True,
momentum=0.1,
fine_tune=False,
# readout args
init_noise=4.1232e-05,
normalize=False,
readout_bias=True,
gamma_readout=0.0076,
share_features=False,
):
"""
Model class of a task-driven transfer core and a factorized (sxf) readout
Args:
dataloaders: a dictionary of dataloaders, one loader per session
in the format {'data_key': dataloader object, .. }
seed: random seed
elu_offset: Offset for the output non-linearity [F.elu(x + self.offset)]
all other args: See Documentation of TransferLearningCore in neuralpredictors.layers.cores and
fullSXF in neuralpredictors.layers.readouts
Returns: An initialized model which consists of model.core and model.readout
"""
if data_info is not None:
n_neurons_dict, in_shapes_dict, input_channels = unpack_data_info(
data_info)
else:
if "train" in dataloaders.keys():
dataloaders = dataloaders["train"]
# Obtain the named tuple fields from the first entry of the first dataloader in the dictionary
in_name, out_name = next(iter(list(dataloaders.values())[0]))._fields
session_shape_dict = get_dims_for_loader_dict(dataloaders)
n_neurons_dict = {
k: v[out_name][1]
for k, v in session_shape_dict.items()
}
in_shapes_dict = {k: v[in_name] for k, v in session_shape_dict.items()}
input_channels = [v[in_name][1] for v in session_shape_dict.values()]
core_input_channels = list(input_channels.values())[0] if isinstance(
input_channels, dict) else input_channels[0]
shared_match_ids = None
if share_features:
shared_match_ids = {
k: v.dataset.neurons.multi_match_id
for k, v in dataloaders.items()
}
all_multi_unit_ids = set(np.hstack(shared_match_ids.values()))
for match_id in shared_match_ids.values():
assert len(set(match_id) & all_multi_unit_ids) == len(
all_multi_unit_ids
), "All multi unit IDs must be present in all datasets"
set_random_seed(seed)
core = TransferLearningCore(
input_channels=core_input_channels,
tl_model_name=tl_model_name,
layers=layers,
pretrained=pretrained,
final_batchnorm=final_batchnorm,
final_nonlinearity=final_nonlinearity,
momentum=momentum,
fine_tune=fine_tune,
)
readout = MultipleFullSXF(
core,
in_shape_dict=in_shapes_dict,
n_neurons_dict=n_neurons_dict,
init_noise=init_noise,
bias=readout_bias,
gamma_readout=gamma_readout,
normalize=normalize,
share_features=share_features,
shared_match_ids=shared_match_ids,
)
# initializing readout bias to mean response
if readout_bias and data_info is None:
for key, value in dataloaders.items():
_, targets = next(iter(value))
readout[key].bias.data = targets.mean(0)
model = Encoder(core, readout, elu_offset)
return model
def taskdriven_fullgaussian2d(
dataloaders,
seed,
elu_offset=0,
data_info=None,
# core args
tl_model_name="vgg16",
layers=4,
pretrained=True,
final_batchnorm=True,
final_nonlinearity=True,
momentum=0.1,
fine_tune=False,
# readout args
init_mu_range=0.3,
init_sigma=0.1,
readout_bias=True,
gamma_readout=0.0076,
gauss_type="full",
grid_mean_predictor={
"type": "cortex",
"input_dimensions": 2,
"hidden_layers": 0,
"hidden_features": 30,
"final_tanh": True,
},
share_features=False,
share_grid=False,
share_transform=False,
init_noise=1e-3,
init_transform_scale=0.2,
):
"""
Model class of a task-driven transfer core and a Gaussian readout
Args:
dataloaders: a dictionary of dataloaders, one loader per session
in the format {'data_key': dataloader object, .. }
seed: random seed
elu_offset: Offset for the output non-linearity [F.elu(x + self.offset)]
grid_mean_predictor: if not None, needs to be a dictionary of the form
{
'type': 'cortex',
'input_dimensions': 2,
'hidden_layers':0,
'hidden_features':20,
'final_tanh': False,
}
In that case the datasets need to have the property `neurons.cell_motor_coordinates`
share_features: whether to share features between readouts. This requires that the datasets
have the properties `neurons.multi_match_id` which are used for matching. Every dataset
has to have all these ids and cannot have any more.
share_grid: whether to share the grid between neurons. This requires that the datasets
have the properties `neurons.multi_match_id` which are used for matching. Every dataset
has to have all these ids and cannot have any more.
share_transform: whether to share the transform from the grid_mean_predictor between neurons. This requires that the datasets
have the properties `neurons.multi_match_id` which are used for matching. Every dataset
has to have all these ids and cannot have any more.
init_noise: noise for initialization of weights
init_transform_scale: scale of the weights of the randomly intialized grid_mean_predictor network
all other args: See Documentation of TransferLearningCore in neuralpredictors.layers.cores and
FullGaussian2d in neuralpredictors.layers.readouts
Returns: An initialized model which consists of model.core and model.readout
"""
if data_info is not None:
n_neurons_dict, in_shapes_dict, input_channels = unpack_data_info(
data_info)
else:
if "train" in dataloaders.keys():
dataloaders = dataloaders["train"]
# Obtain the named tuple fields from the first entry of the first dataloader in the dictionary
in_name, out_name = next(iter(list(dataloaders.values())[0]))._fields
session_shape_dict = get_dims_for_loader_dict(dataloaders)
n_neurons_dict = {
k: v[out_name][1]
for k, v in session_shape_dict.items()
}
in_shapes_dict = {k: v[in_name] for k, v in session_shape_dict.items()}
input_channels = [v[in_name][1] for v in session_shape_dict.values()]
core_input_channels = list(input_channels.values())[0] if isinstance(
input_channels, dict) else input_channels[0]
source_grids = None
grid_mean_predictor_type = None
if grid_mean_predictor is not None:
grid_mean_predictor = copy.deepcopy(grid_mean_predictor)
grid_mean_predictor_type = grid_mean_predictor.pop("type")
if grid_mean_predictor_type == "cortex":
input_dim = grid_mean_predictor.pop("input_dimensions", 2)
source_grids = {}
for k, v in dataloaders.items():
# real data
try:
if v.dataset.neurons.animal_ids[0] != 0:
source_grids[
k] = v.dataset.neurons.cell_motor_coordinates[:, :
input_dim]
# simulated data -> get random linear non-degenerate transform of true positions
else:
source_grid_true = v.dataset.neurons.center[:, :
input_dim]
det = 0.0
loops = 0
grid_bias = np.random.rand(2) * 3
while det < 5.0 and loops < 100:
matrix = np.random.rand(2, 2) * 3
det = np.linalg.det(matrix)
loops += 1
assert det > 5.0, "Did not find a non-degenerate matrix"
source_grids[k] = np.add(
(matrix @ source_grid_true.T).T, grid_bias)
except FileNotFoundError:
print(
"Dataset type is not recognized to be from Baylor College of Medicine."
)
source_grids[
k] = v.dataset.neurons.cell_motor_coordinates[:, :
input_dim]
elif grid_mean_predictor_type == "shared":
pass
else:
raise ValueError(
"Grid mean predictor type {} not understood.".format(
grid_mean_predictor_type))
shared_match_ids = None
if share_features or share_grid:
shared_match_ids = {
k: v.dataset.neurons.multi_match_id
for k, v in dataloaders.items()
}
all_multi_unit_ids = set(np.hstack(shared_match_ids.values()))
for match_id in shared_match_ids.values():
assert len(set(match_id) & all_multi_unit_ids) == len(
all_multi_unit_ids
), "All multi unit IDs must be present in all datasets"
set_random_seed(seed)
core = TransferLearningCore(
input_channels=core_input_channels,
tl_model_name=tl_model_name,
layers=layers,
pretrained=pretrained,
final_batchnorm=final_batchnorm,
final_nonlinearity=final_nonlinearity,
momentum=momentum,
fine_tune=fine_tune,
)
readout = MultipleFullGaussian2d(
core,
in_shape_dict=in_shapes_dict,
n_neurons_dict=n_neurons_dict,
init_mu_range=init_mu_range,
bias=readout_bias,
init_sigma=init_sigma,
gamma_readout=gamma_readout,
gauss_type=gauss_type,
grid_mean_predictor=grid_mean_predictor,
grid_mean_predictor_type=grid_mean_predictor_type,
source_grids=source_grids,
share_features=share_features,
share_grid=share_grid,
shared_match_ids=shared_match_ids,
share_transform=share_transform,
init_noise=init_noise,
init_transform_scale=init_transform_scale,
)
# initializing readout bias to mean response
if readout_bias and data_info is None:
for key, value in dataloaders.items():
_, targets = next(iter(value))
readout[key].bias.data = targets.mean(0)
model = Encoder(core, readout, elu_offset)
return model
def se2d_fullSXF(
dataloaders,
seed,
elu_offset=0,
data_info=None,
transfer_state_dict=None,
# core args
hidden_channels=64,
input_kern=9,
hidden_kern=7,
layers=4,
gamma_input=6.3831,
skip=0,
bias=False,
final_nonlinearity=True,
momentum=0.9,
pad_input=False,
batch_norm=True,
hidden_dilation=1,
laplace_padding=None,
input_regularizer="LaplaceL2norm",
stack=-1,
se_reduction=32,
n_se_blocks=0,
depth_separable=True,
linear=False,
init_noise=4.1232e-05,
normalize=False,
readout_bias=True,
gamma_readout=0.0076,
share_features=False,
):
"""
Model class of a SE2dCore and a factorized (sxf) readout
Args:
dataloaders: a dictionary of dataloaders, one loader per session
in the format {'data_key': dataloader object, .. }
seed: random seed
elu_offset: Offset for the output non-linearity [F.elu(x + self.offset)]
all other args: See Documentation of SE2dCore in neuralpredictors.layers.cores and
fullSXF in neuralpredictors.layers.readouts
Returns: An initialized model which consists of model.core and model.readout
"""
if transfer_state_dict is not None:
print(
"Transfer state_dict given. This will only have an effect in the bayesian hypersearch. See: TrainedModelBayesianTransfer "
)
if data_info is not None:
n_neurons_dict, in_shapes_dict, input_channels = unpack_data_info(
data_info)
else:
if "train" in dataloaders.keys():
dataloaders = dataloaders["train"]
# Obtain the named tuple fields from the first entry of the first dataloader in the dictionary
in_name, out_name = next(iter(list(dataloaders.values())[0]))._fields
session_shape_dict = get_dims_for_loader_dict(dataloaders)
n_neurons_dict = {
k: v[out_name][1]
for k, v in session_shape_dict.items()
}
in_shapes_dict = {k: v[in_name] for k, v in session_shape_dict.items()}
input_channels = [v[in_name][1] for v in session_shape_dict.values()]
core_input_channels = list(input_channels.values())[0] if isinstance(
input_channels, dict) else input_channels[0]
shared_match_ids = None
if share_features:
shared_match_ids = {
k: v.dataset.neurons.multi_match_id
for k, v in dataloaders.items()
}
all_multi_unit_ids = set(np.hstack(shared_match_ids.values()))
for match_id in shared_match_ids.values():
assert len(set(match_id) & all_multi_unit_ids) == len(
all_multi_unit_ids
), "All multi unit IDs must be present in all datasets"
set_random_seed(seed)
core = SE2dCore(
input_channels=core_input_channels,
hidden_channels=hidden_channels,
input_kern=input_kern,
hidden_kern=hidden_kern,
layers=layers,
gamma_input=gamma_input,
skip=skip,
final_nonlinearity=final_nonlinearity,
bias=bias,
momentum=momentum,
pad_input=pad_input,
batch_norm=batch_norm,
hidden_dilation=hidden_dilation,
laplace_padding=laplace_padding,
input_regularizer=input_regularizer,
stack=stack,
se_reduction=se_reduction,
n_se_blocks=n_se_blocks,
depth_separable=depth_separable,
linear=linear,
)
readout = MultipleFullSXF(
core,
in_shape_dict=in_shapes_dict,
n_neurons_dict=n_neurons_dict,
init_noise=init_noise,
bias=readout_bias,
gamma_readout=gamma_readout,
normalize=normalize,
share_features=share_features,
shared_match_ids=shared_match_ids,
)
# initializing readout bias to mean response
if readout_bias and data_info is None:
for key, value in dataloaders.items():
_, targets = next(iter(value))
readout[key].bias.data = targets.mean(0)
model = Encoder(core, readout, elu_offset)
return model
def se2d_pointpooled(
dataloaders,
seed,
elu_offset=0,
data_info=None,
# core args
hidden_channels=64,
input_kern=9, # core args
hidden_kern=7,
layers=4,
gamma_input=46.402,
bias=False,
skip=0,
final_nonlinearity=True,
momentum=0.9,
pad_input=False,
batch_norm=True,
hidden_dilation=1,
laplace_padding=None,
input_regularizer="LaplaceL2norm",
stack=-1,
se_reduction=32,
n_se_blocks=0,
depth_separable=True,
linear=False,
# readout args
pool_steps=2,
pool_kern=3,
readout_bias=True,
gamma_readout=0.0207,
init_range=0.2,
):
"""
Model class of a SE2dCore and a pointpooled (spatial transformer) readout
Args:
dataloaders: a dictionary of dataloaders, one loader per session
in the format {'data_key': dataloader object, .. }
seed: random seed
elu_offset: Offset for the output non-linearity [F.elu(x + self.offset)]
all other args: See Documentation of SE2dCore in neuralpredictors.layers.cores and
PointPooled2D in neuralpredictors.layers.readouts
Returns: An initialized model which consists of model.core and model.readout
"""
if data_info is not None:
n_neurons_dict, in_shapes_dict, input_channels = unpack_data_info(
data_info)
else:
if "train" in dataloaders.keys():
dataloaders = dataloaders["train"]
# Obtain the named tuple fields from the first entry of the first dataloader in the dictionary
in_name, out_name = next(iter(list(dataloaders.values())[0]))._fields
session_shape_dict = get_dims_for_loader_dict(dataloaders)
n_neurons_dict = {
k: v[out_name][1]
for k, v in session_shape_dict.items()
}
in_shapes_dict = {k: v[in_name] for k, v in session_shape_dict.items()}
input_channels = [v[in_name][1] for v in session_shape_dict.values()]
core_input_channels = list(input_channels.values())[0] if isinstance(
input_channels, dict) else input_channels[0]
set_random_seed(seed)
core = SE2dCore(
input_channels=core_input_channels,
hidden_channels=hidden_channels,
input_kern=input_kern,
hidden_kern=hidden_kern,
layers=layers,
gamma_input=gamma_input,
bias=bias,
skip=skip,
final_nonlinearity=final_nonlinearity,
momentum=momentum,
pad_input=pad_input,
batch_norm=batch_norm,
hidden_dilation=hidden_dilation,
laplace_padding=laplace_padding,
input_regularizer=input_regularizer,
stack=stack,
se_reduction=se_reduction,
n_se_blocks=n_se_blocks,
depth_separable=depth_separable,
linear=linear,
)
readout = MultiplePointPooled2d(
core,
in_shape_dict=in_shapes_dict,
n_neurons_dict=n_neurons_dict,
pool_steps=pool_steps,
pool_kern=pool_kern,
bias=readout_bias,
gamma_readout=gamma_readout,
init_range=init_range,
)
# initializing readout bias to mean response
if readout_bias and data_info is None:
for key, value in dataloaders.items():
_, targets = next(iter(value))
readout[key].bias.data = targets.mean(0)
model = Encoder(core, readout, elu_offset)
return model