<jupyter_start><jupyter_text>Parameters<jupyter_code># Data params
input_dim = 784
output_dim = 10

# Network params
n_hidden_units = 2000
activation_fn = tf.nn.relu

# Optimization params
batch_size = 256
epochs_per_task = 20 
learning_rate=1e-3
xi = 0.1

# Reset optimizer after each age
reset_optimizer = False<jupyter_output><empty_output><jupyter_text>Construct datasets<jupyter_code>n_tasks = 10
full_datasets, final_test_datasets = utils.construct_permute_mnist(num_tasks=n_tasks)
# training_datasets, validation_datasets = utils.mk_training_validation_splits(full_datasets, split_fractions=(0.9, 0.1))
training_datasets = full_datasets
validation_datasets = final_test_datasets<jupyter_output><empty_output><jupyter_text>Construct network, loss, and updates<jupyter_code>tf.reset_default_graph()
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.InteractiveSession(config=config)
sess.run(tf.global_variables_initializer())
from keras.models import Sequential
from keras.layers import Dense, Dropout
model = Sequential()
model.add(Dense(n_hidden_units, activation=activation_fn, input_dim=input_dim))
model.add(Dense(n_hidden_units, activation=activation_fn))
model.add(Dense(output_dim, activation='softmax'))
from pathint import protocols
from pathint.optimizers import KOOptimizer
from keras.optimizers import SGD, Adam, RMSprop
from keras.callbacks import Callback
from pathint.keras_utils import LossHistory


protocol_name, protocol = protocols.PATH_INT_PROTOCOL(omega_decay='sum', xi=xi)
# protocol_name, protocol = protocols.FISHER_PROTOCOL(omega_decay='sum')
opt = Adam(lr=learning_rate, beta_1=0.9, beta_2=0.999)
opt_name = 'adam'
# opt = SGD(lr=learning_rate)
# opt_name = 'sgd'
oopt = KOOptimizer(opt, model=model, **protocol)
model.compile(loss="categorical_crossentropy", optimizer=oopt, metrics=['accuracy'])

history = LossHistory()
callbacks = [history]
file_prefix = "data_%s_opt%s_lr%.2e_bs%i_ep%i_tsks%i"%(protocol_name, opt_name, learning_rate, batch_size, epochs_per_task, n_tasks)
datafile_name = "%s.pkl.gz"%(file_prefix)<jupyter_output><empty_output><jupyter_text>Train!<jupyter_code>diag_vals = dict()
all_evals = dict()
# all_evals = utils.load_zipped_pickle(datafile_name)
# returns empty dict if file not found

def run_fits(cvals, training_data, valid_data, eval_on_train_set=False):
    for cidx, cval_ in enumerate(cvals):
        fs = []
        evals = []
        sess.run(tf.global_variables_initializer())
        cstuffs = []
        cval = cval_
        print( "setting cval")
        oopt.set_strength(cval)
        print("cval is %e"%sess.run(oopt.lam))
        for age, tidx in enumerate(range(n_tasks)):
            print("Age %i, cval is=%f"%(age,cval))
            oopt.set_nb_data(len(training_data[tidx][0]))
            stuffs = model.fit(training_data[tidx][0], training_data[tidx][1], batch_size, epochs_per_task, callbacks=callbacks,
                              verbose=0)
            oopt.update_task_metrics(training_data[tidx][0], training_data[tidx][1], batch_size)
            oopt.update_task_vars()
            ftask = []
            for j in range(n_tasks):
                if eval_on_train_set:
                    f_ = model.evaluate(training_data[j][0], training_data[j][1], batch_size, verbose=0)
                else:
                    f_ = model.evaluate(valid_data[j][0], valid_data[j][1], batch_size, verbose=0)
                ftask.append(np.mean(f_[1]))
            evals.append(ftask)
            cstuffs.append(stuffs)
            
            # Re-initialize optimizer variables
            if reset_optimizer:
                oopt.reset_optimizer()

        # diag_vals[cval_] = oopt.get_numvals('omega')
        evals = np.array(evals)
        all_evals[cval_] = evals
        
        # backup all_evals to disk
        utils.save_zipped_pickle(all_evals, datafile_name)
cvals = [0, 0.01, 0.1, 1.0]
print(cvals)
# %%capture
run_fits(cvals, training_datasets, validation_datasets)
# backup all_evals to disk
# all_evals = dict() # uncomment to delete on disk
utils.save_zipped_pickle(all_evals, datafile_name)
o = oopt.get_numvals_list('omega')
cmap = plt.get_cmap('cool') 
cNorm  = colors.Normalize(vmin=-4, vmax=np.log(np.max(list(all_evals.keys()))))
# cNorm  = colors.Normalize()
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cmap)
print(scalarMap.get_clim())
figure(figsize=(14, 4))
axs = [subplot(1,n_tasks+1,1)]#, None, None]
for i in range(1, n_tasks + 1):
    axs.append(subplot(1, n_tasks+1, i+1, sharex=axs[0], sharey=axs[0]))
    
keys = list(all_evals.keys())
sorted_keys = np.sort(keys)

for cval in sorted_keys:
    evals = all_evals[cval]
    for j in range(n_tasks):
        colorVal = scalarMap.to_rgba(np.log(cval))
        axs[j].plot(evals[:, j], c=colorVal)#, label="t%d, c%g"%(j, cval))
    label = "c=%g"%cval
    average = evals.mean(1)
    axs[-1].plot(average, c=colorVal, label=label)
    
for i, ax in enumerate(axs):
    ax.legend(bbox_to_anchor=(1.0,1.0))
    ax.set_title((['task %d'%j for j in range(n_tasks)] + ['average'])[i])
gcf().tight_layout()
for cval in sorted_keys:
    stuff = []
    for i in range(len(all_evals[cval])):#n_tasks):
        stuff.append(all_evals[cval][i][:i+1].mean())
    plot(range(1,n_tasks+1), stuff, 'o-', label="c=%g"%cval)
axhline(all_evals[cval][0][0], linestyle='--', color='k')
    
xlabel('Number of tasks')
ylabel('Fraction correct')
legend(loc='best')
ylim(0.9, 1.02)
xlim(0.5, 10.5)
grid('on')

savefig("%s.pdf"%(file_prefix))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13553510196715393, 'fraction_numerical': 0.016062082656560187, 'mean_word_length': 4.422700587084149, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '491075', 'license': 'MIT', 'n_tokens_mistral': 2118, 'n_tokens_neox': 2053, 'n_words': 350, 'path': 'fig_permuted_mnist/Permuted MNIST.ipynb', 'repo_name': 'ganguli-lab/pathint', 'size': '5541'}