<jupyter_start><jupyter_text>Load MNIST<jupyter_code>from torchvision import datasets, transforms

n_train = 4000

train_set = datasets.MNIST('/mnt/ssd3/ronan/pytorch', train=True, download=True)
X_train = train_set.data.numpy()[:n_train]
X_train = X_train.reshape((n_train, -1))
y_train = train_set.targets.numpy()[:n_train]
del train_set

test_set = datasets.MNIST('/mnt/ssd3/ronan/pytorch', train=False, download=True)
X_test = test_set.data.numpy()
X_test = X_test.reshape((X_test.shape[0], -1))
y_test = test_set.targets.numpy()
n_hidden = get_units_at_interp(n_train, k=len(np.unique(y_train)), d=X_train.shape[-1])
print(f'{n_hidden} hidden units')
results = [[
    'hidden_layer_dims',
    'train_01_error',
    'train_mse',
    'test_01_error',
    'test_mse',   
] + [
    'IRM_L1', 'IRM_L2', 'n_regions', 'ACTS_L2', 'IRM_h*', 'ACTS_h*', 'entropy',
    'rows_mean_L2', 'cols_mean_L1', 'cols_mean_L2'
] + [
    'n_parameters', 'depth', 'width'
]]

NETWORK_PARAMS = {
    'hidden_layer_dims': [
        [4],
        [8],
        [16],
        [32],
        [40],
        [45],
        [47],
        [51],
        [55],
        [60],
        [100]
    ],
    'n_epochs': [500],# [0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024],# 2048],
    'learning_rate': [1e-5],
    'batch_size': [128],
    'verbose': [0],
    'early_stop_thresh': [0],
    "bias": [True],
    # "loss": [torch.nn.MSELoss],
}
keys, values = zip(*NETWORK_PARAMS.items())
model_params_grid = [dict(zip(keys, v)) for v in itertools.product(*values)]

model = None
for model_params in tqdm(model_params_grid):
    # Train and test model
    model, y_train_pred, y_test_pred, model_metrics = run_network(
        X_train, y_train, X_test, model_params, model
    )

    # Compute metrics
    results += [
        [model_params['hidden_layer_dims'][0]]
        + get_y_metrics(y_train, y_train_pred)  # Train
        + get_y_metrics(y_test, y_test_pred)  # Test
        + model_metrics
    ]
df = pd.DataFrame(results[1:], columns=results[0])
df
train_error = 'train_01_error'
test_error = 'test_01_error'

# Create new metrics/columns
df['generalization_gap'] = df[test_error] - df[train_error]
sns.set_style('white')
sns.set_context('notebook')

fig, ax = plt.subplots(1, 1, sharex=True)

sns.regplot(data=df, x='n_parameters', y=test_error, lowess=False, ax=ax, label='test', fit_reg=False)
sns.regplot(data=df, x='n_parameters', y=train_error, lowess=False, ax=ax, label='train', fit_reg=False)
# sns.regplot(data=df, x='n_parameters', y='generalization_gap', lowess=False, ax=ax, label='gen_gap')

ax.axvline(len(np.unique(y_train))*n_train, ls='--', c='black')
plt.xscale('log')
# plt.yscale('symlog')
plt.legend()
plt.show()
corr_vars = [
    'generalization_gap',
    'n_parameters', 'depth', 'width', # 'n_epochs',
    'IRM_L1', 'IRM_L2', 'n_regions', 'ACTS_L2',
    'IRM_h*', 'ACTS_h*', 'entropy',
    'rows_mean_L2',
    'cols_mean_L1', 'cols_mean_L2',
]

corr_df = df[corr_vars].corr('kendall')
display(corr_df)<jupyter_output><empty_output><jupyter_text>XOR<jupyter_code>n_train = 1024

X_train, y_train = generate_gaussian_parity(
    n_samples=n_train, angle_params=0, random_state=0, cov_scale=1
)

X_test, y_test = generate_gaussian_parity(
    n_samples=10000, angle_params=0, random_state=1234, cov_scale=1
)

import matplotlib.pyplot as plt
plt.scatter(X_train[:, 0], X_train[:, 1], c=y_train)
plt.show()
n_hidden = get_units_at_interp(n_train, k=len(np.unique(y_train)), d=X_train.shape[-1])
print(f'{n_hidden} hidden units')
results = [[
    'hidden_layer_dims',
    'train_01_error',
    'train_mse',
    'test_01_error',
    'test_mse',   
] + [
    'IRM_L1', 'IRM_L2', 'n_regions', 'ACTS_L2', 'IRM_h*', 'ACTS_h*', 'entropy',
    'rows_mean_L2', 'cols_mean_L1', 'cols_mean_L2'
] + [
    'n_parameters', 'depth', 'width'
]]

NETWORK_PARAMS = {
    'hidden_layer_dims': [
        [100],
        [300],
        [450],
        [500],
        [512],
        [550],
        [600],
        [1000]
    ],
    'n_epochs': [1000],# [0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024],# 2048],
    'learning_rate': [1e-5],
    'batch_size': [128],
    'verbose': [0],
    'early_stop_thresh': [0],
    "bias": [True],
    # "loss": [torch.nn.MSELoss],
}
keys, values = zip(*NETWORK_PARAMS.items())
model_params_grid = [dict(zip(keys, v)) for v in itertools.product(*values)]

model = None
for model_params in tqdm(model_params_grid):
    # Train and test model
    model, y_train_pred, y_test_pred, model_metrics = run_network(
        X_train, y_train, X_test, model_params, model
    )

    # Compute metrics
    results += [
        [model_params['hidden_layer_dims'][0]]
        + get_y_metrics(y_train, y_train_pred)  # Train
        + get_y_metrics(y_test, y_test_pred)  # Test
        + model_metrics
    ]
df = pd.DataFrame(results[1:], columns=results[0])
train_error = 'train_01_error'
test_error = 'test_01_error'

# Create new metrics/columns
df['generalization_gap'] = df[test_error] - df[train_error]

sns.set_style('white')
sns.set_context('notebook')

fig, ax = plt.subplots(1, 1, sharex=True)

sns.regplot(data=df, x='n_parameters', y=test_error, lowess=False, ax=ax, label='test', fit_reg=False)
sns.regplot(data=df, x='n_parameters', y=train_error, lowess=False, ax=ax, label='train', fit_reg=False)
# sns.regplot(data=df, x='n_parameters', y='generalization_gap', lowess=False, ax=ax, label='gen_gap')

ax.axvline(len(np.unique(y_train))*n_train, ls='--', c='black')
plt.xscale('log')
# plt.yscale('symlog')
plt.legend()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17544490382886932, 'fraction_numerical': 0.035592306309545206, 'mean_word_length': 4.360308285163777, 'pattern_counts': {'":': 4, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '550843', 'license': 'Apache-2.0', 'n_tokens_mistral': 2425, 'n_tokens_neox': 2293, 'n_words': 382, 'path': 'notebooks/belkin_replication.ipynb', 'repo_name': 'rflperry/double_descent', 'size': '5563'}