<jupyter_start><jupyter_text>Training ConfigurationConfiguration of the model, the trainer and the dataset used for training.<jupyter_code>scan = '20457-5-9'
paths = [f'data/static{scan}-preproc0/']


model_fn = 'model_components.se_core_spatialXfeature_readout'
model_config = {
   'pad_input': False,
   'stack': -1,
   'layers':4,
   'input_kern': 15,
   'gamma_input': 1., #20 #6.3831,          
   'gamma_readout':  0.002362354239446914, #0.0076,
   'hidden_dilation': 1,
   'hidden_kern': 13,
   'hidden_channels': 64,
   'n_se_blocks': 0,
   'depth_separable': True,
   'normalize': False,
    'init_noise' :4.1232e-05,
}

trainer_fn = 'cotton2020.nnvision_trainer'
trainer_config = dict(max_iter=100,
                      verbose=False, 
                      lr_decay_steps=4,
                      avg_loss=False, 
                      patience=5,
                      lr_init=.0041)

readout_key = f'{scan}-0/'<jupyter_output><empty_output><jupyter_text>Train Full Klindt et al. model  Get DataLoader<jupyter_code>unit_ids = np.load(f'configs/{scan}_train_units.npy')

dataset_fn = 'cotton2020.neurips_loaders'
dataset_config = dict(
    paths=paths,
    areas = ('V1',),
    batch_size=64,
    normalize=True,
    neuron_ids=[unit_ids],
)

dataloaders = builder.get_data(dataset_fn, dataset_config)
dataloaders<jupyter_output><empty_output><jupyter_text>Define Model<jupyter_code>model = builder.get_model(model_fn=model_fn, 
                          model_config=model_config, 
                          dataloaders=dataloaders, 
                          seed=1000)
model<jupyter_output><empty_output><jupyter_text>Get Trainer<jupyter_code>trainer = builder.get_trainer(trainer_fn, trainer_config)<jupyter_output><empty_output><jupyter_text>Run Training<jupyter_code>score, output, model_state = trainer(model=model, dataloaders=dataloaders, seed=1000)
state_dict = model.state_dict()
with open(f'models/{scan}_full_model_statedict.pkl', 'wb') as fid:
    pickle.dump(state_dict, fid)<jupyter_output><empty_output><jupyter_text>Train Readout only<jupyter_code>total_results = dict()
data_size = 1000

unit_ids = np.load(f'configs/{scan}_test_units.npy')
train_trials =  np.load(f'configs/{scan}_train_trials.npy')
dataset_fn = 'cotton2020.neurips_loaders'
dataset_config = dict(
    paths=paths,
    batch_size=64,
    normalize=True,
    neuron_ids=[unit_ids],
    trial_ids = dict(train=train_trials[:data_size]) if data_size <= 1000 else None
)
dataloaders = builder.get_data(dataset_fn, dataset_config)
dataloaders
model = builder.get_model(model_fn=model_fn, 
                          model_config=model_config, 
                          dataloaders=dataloaders, 
                          seed=1000)
model<jupyter_output><empty_output><jupyter_text>Load old model, copy core, and freeze it<jupyter_code>with open(f'models/{scan}_full_model_statedict.pkl', 'rb') as fid:
    state_dict = pickle.load(fid)
    core_state_dict = OrderedDict(
        [(k[5:], v) for k,v in state_dict.items() if k.startswith('core.')]
    )
old_state = OrderedDict(core_state_dict) 
model.core.load_state_dict(core_state_dict)


for param in model.core.parameters():
    param.requires_grad = False<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>trainer = builder.get_trainer(trainer_fn, trainer_config)
score, output, model_state = trainer(model=model, dataloaders=dataloaders, seed=1000)<jupyter_output>Epoch 1:  62%|██████▎   | 10/16 [00:00<00:00, 10.66it/s]<jupyter_text>Make sure core was actually frozen<jupyter_code>for k, v in model.core.state_dict().items():
    if not 'running_mean' in k and not 'running_var' in k and not 'num_batches_tracked' in k:
        print(f'Checking {k}', (v - old_state[k]).abs().max())
        assert (v - old_state[k]).abs().max() < 1e-6<jupyter_output><empty_output><jupyter_text>Save model<jupyter_code>state_dict = model.state_dict()
with open(f'models/{scan}_datasize{data_size}readout_model_statedict.pkl', 'wb') as fid:
    pickle.dump(state_dict, fid)<jupyter_output><empty_output><jupyter_text>Evaluate model<jupyter_code>corrs = get_correlations(model, dataloaders['test'])
oracles = get_oracles(dataloaders["test"], as_dict=True)
corrs.mean()
a, _, _, _ = np.linalg.lstsq(oracles[readout_key][:, None], corrs)
print(f'Percent oracle is {a[0]*100:.2f}%')
total_results[data_size] =  dict(
    unit_ids = dataloaders['train'][readout_key].dataset.neurons.unit_ids,
    test_correlation = corrs.squeeze(),
    oracle_correlation = oracles[readout_key],
    mean_test_correlation = corrs.mean(),
    fraction_oracle = a
)
np.savez(f'results/{scan}-data_efficieny.npz', total_results)
{k:total_results[k]['mean_test_correlation'] for k in total_results}<jupyter_output><empty_output><jupyter_text>Values from the Cotton et al. 2020<jupyter_code>{50: 0.046387184,
 100: 0.07622242,
 250: 0.12529792,
 500: 0.17531049,
 750: 0.19878243,
 1000: 0.21516044,
 4994: 0.26639703}<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1288244766505636, 'fraction_numerical': 0.04528985507246377, 'mean_word_length': 4.646590909090909, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 48, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '553084', 'license': 'Apache-2.0', 'n_tokens_mistral': 2040, 'n_tokens_neox': 1850, 'n_words': 277, 'path': 'pno/notebooks/Cotton2020.ipynb', 'repo_name': 'sinzlab/fnp_neurips2020', 'size': '4968'}