import os
from os.path import join as osjoin
from pathlib import Path
from copy import deepcopy
from functools import partial
import argparse
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from joblib import Parallel, delayed
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchaudio
import torchaudio.datasets as tads
import torchvision.datasets as tvds
import torchvision.transforms as transforms
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import hetero_utils as hu
datapath = torch.hub.get_dir()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(torch.__version__)
print(torchaudio.__version__)
print(f'device is: {device}')
print(f'data path is {datapath}')
# DEMO_NAME = 'machlearn'
# SAVE_PATH = osjoin(Path(__file__).resolve().parents[0], DEMO_NAME)
SAVE_PATH = osjoin(Path(__file__).resolve().parents[0], 'results')
Path(SAVE_PATH).mkdir(parents=True, exist_ok=True)
INTEGRATION_SETTINGS = {
'speechcommands': dict(reset_dim=0),
'librespeech' : dict(reset_dim=0),
'mnist' : dict(reset_dim=0, u_reset=-1),
}
# Taken from https://arxiv.org/pdf/1804.03209
SPEECHCOMMAND_LABEL_COUNTS = {
"Backward": 1664,
"Bed": 2014,
"Bird": 2064,
"Cat": 2031,
"Dog": 2128,
"Down": 3917,
"Eight": 3787,
"Five": 4052,
"Follow": 1579,
"Forward": 1557,
"Four": 3728,
"Go": 3880,
"Happy": 2054,
"House": 2113,
"Learn": 1575,
"Left": 3801,
"Marvin": 2100,
"Nine": 3934,
"No": 3941,
"Off": 3745,
"On": 3845,
"One": 3890,
"Right": 3778,
"Seven": 3998,
"Sheila": 2022,
"Six": 3860,
"Stop": 3872,
"Three": 3727,
"Tree": 1759,
"Two": 3880,
"Up": 3723,
"Visual": 1592,
"Wow": 2123,
"Yes": 4044,
"Zero": 4052,
}
[docs]
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
[docs]
def get_dataset(name='speechcommands', subset='train'):
if name=='librespeech':
if subset == 'train':
return tads.LIBRISPEECH(datapath, url="train-clean-100", download=True)
else:
return tads.LIBRISPEECH(datapath, url="test-clean", download=True)
elif name == 'speechcommands':
if subset == 'train':
return tads.SPEECHCOMMANDS(root=datapath, download=True, subset='training',)
else:
return tads.SPEECHCOMMANDS(root=datapath, download=True, subset='testing')
elif name == 'mnist':
if subset == 'train':
return tvds.MNIST(root=datapath, download=True, train=True, transform=transforms.ToTensor())
else:
return tvds.MNIST(root=datapath, download=True, train=False, transform=transforms.ToTensor())
[docs]
def get_bundle(name = 'emformer',
return_feature_extractor=True,
return_decoder=False,
return_token_processor=False):
if name == 'emformer':
bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH
else:
raise NotImplementedError
misc = []
if return_feature_extractor:
misc.append(bundle.get_feature_extractor())
if return_decoder:
misc.append(bundle.get_decoder())
if return_token_processor:
misc.append(bundle.get_token_processor())
return bundle, misc
[docs]
def collate_fn_mnist(batch, stream=True, permutation=None):
images, labels = zip(*batch)
pixels = torch.stack(images).numpy().squeeze() # shape: (B, 28, 28)
pixels = pixels.reshape(-1, 784)
if permutation is not None:
pixels = pixels[:, permutation]
# unit-norm
# pixels /= torch.linalg.norm(pixels)
if stream:
# prepend -1 to the pixels (timestamps) as the "fixation/reset cue"
pixels = np.hstack((-1*np.ones(len(batch))[:,np.newaxis], pixels))
sample_length = pixels.shape[1]
if stream:
pixels = np.expand_dims(pixels, 1)
pixels = pixels.reshape(-1, 1)
labels = np.array(labels)
return pixels, sample_length*np.ones(len(batch), dtype=int), labels
[docs]
def collate_fn_arged(batch, feature_extractor, stream, keep_classes):
features, lens, rates, targets = [], [], [], []
# print(keep_classes)
for waveform, rate, label, spkr, utter in batch:
# print(label)
if keep_classes == 'all' or label in keep_classes:
# squeeze out channel
waveform = waveform.squeeze()
if not stream:
# pad to the fix length of 16k (1 second with sampling rate of 16 kHz)
# here, pad = (padding_left, padding_right, padding_top, padding_bottom)
waveform = F.pad(waveform, pad=(0, 16000-waveform.shape[0]))
# Pad sequences to the max length in this specific batch
# This makes them (batch, time, channels)
# tensors = torch.nn.utils.rnn.pad_sequence(tensors, batch_first=True, padding_value=0.0)
# max-normalize
waveform /= waveform.abs().max()
# exctract acoustic features
feature, feature_len = feature_extractor(waveform)
# add the reset-cue (fixation)
if stream:
fix = torch.zeros(feature_len)
fix[0] = 1.
feature = torch.column_stack((fix, feature))
features += [feature]
lens += [feature_len]
rates += [rate]
targets += [label]
# spkrs += [torch.tensor(spkr)]
if stream:
features = torch.concat(features)
else:
features = torch.stack(features)
lens = torch.concat(lens)
# targets = targets
# rates = rates
return features.numpy(), lens.numpy(), targets
[docs]
def get_train_test(dataset_name='speechcommands', seed=1234,
stream=True, Nbatch = 100, batch_size=32,
top_n = 10, permute=True,
center=False, unit_std=False, normed_trial=True,
orthgonalize=False, whiten=True, zero_phase=True):
print(dataset_name, top_n)
rng_cpu = torch.Generator()
rng_cpu.manual_seed(seed)
dataset_trn = get_dataset(dataset_name, 'train')
dataset_tst = get_dataset(dataset_name, 'test')
if dataset_name == 'mnist':
permutation = None
if permute:
rng = np.random.default_rng(seed=seed)
permutation = rng.permutation(784)
collate_fn = partial(collate_fn_mnist,
stream=stream, permutation=permutation)
else:
bundle, misc = get_bundle('emformer')
feature_extractor = misc[0]
if top_n == -1:
keep_classes = 'all'
else:
# the stats taken from https://arxiv.org/pdf/1804.03209
lc = pd.Series(SPEECHCOMMAND_LABEL_COUNTS)
keep_classes = lc.sort_values(ascending=False).index[:top_n]
keep_classes = keep_classes.map(lambda x: x.lower())
keep_classes = list(keep_classes.values)
print ('maintaining the following labels from SC dataset')
print(keep_classes)
collate_fn = partial(collate_fn_arged,
stream=stream,
feature_extractor=feature_extractor,
keep_classes=keep_classes)
dataloader_settings = dict(batch_size=batch_size, shuffle=True,
collate_fn=collate_fn, generator=rng_cpu)
dataloader_trn = DataLoader(dataset_trn, **dataloader_settings)
dataloader_tst = DataLoader(dataset_tst, **dataloader_settings)
trn_generator = iter(dataloader_trn)
tst_generator = iter(dataloader_tst)
us_trn = []
ys_trn = []
ln_trn = []
us_tst = []
ys_tst = []
ln_tst = []
for batch in range(Nbatch):
if batch%20==0:
print(batch)
features, lens, labels = next(trn_generator)
us_trn.append(features)
ys_trn.append(labels)
ln_trn.append(lens)
features, lens, labels = next(tst_generator)
us_tst.append(features)
ys_tst.append(labels)
ln_tst.append(lens)
us_trn = np.concatenate(us_trn)
us_tst = np.concatenate(us_tst)
ys_trn = np.concatenate(ys_trn)
ys_tst = np.concatenate(ys_tst)
ln_trn = np.concatenate(ln_trn)
ln_tst = np.concatenate(ln_tst)
if center:
ss = StandardScaler(with_std=unit_std)
us_trn = ss.fit_transform(us_trn)
us_tst = ss.transform(us_tst)
if orthgonalize:
pca = PCA(whiten=whiten)
trn_t = pca.fit_transform(us_trn[:,1:])
tst_t = pca.transform(us_tst[:, 1:])
if zero_phase:
trn_t = trn_t@pca.components_
tst_t = tst_t@pca.components_
us_trn = np.c_[us_trn[:,0], trn_t]
us_tst = np.c_[us_tst[:,0], tst_t]
if normed_trial:
ons, offs = hu.get_trial_start_end(ln_trn)
for on, off in zip(ons, offs):
us_trn[on:off, 1:] /= np.linalg.norm(us_trn[on:off, 1:])
ons, offs = hu.get_trial_start_end(ln_tst)
for on, off in zip(ons, offs):
us_tst[on:off, 1:] /= np.linalg.norm(us_tst[on:off, 1:])
return (us_trn, ys_trn, ln_trn), (us_tst, ys_tst, ln_tst)
[docs]
def run_single_net(h=np.sqrt(10), N=10, p=.1, seed=1234,
Nbatch=50, dt=1,
unit_std=False, center=False,
orthgonalize=True, zero_phase=True,
filt=False, first_unstream=True, flattening_method='flat',
af = 'tanh', num_snapshot = 10,
plot=True,
):
trn, tst = get_train_test(dataset_name='speechcommands', Nbatch=Nbatch,
unit_std=unit_std, center=center,
orthgonalize=orthgonalize, zero_phase=zero_phase)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
# extract tau
tau0, _ = hu.extract_base_tau(us_trn, dt=dt,
sample_lengths=lens_trn,
filt=filt, first_unstream=first_unstream,
flattening_method=flattening_method, center=False,
savepath=SAVE_PATH, savename='taus'
)
# defining settings
tau_settings = dict(size=N, mu=tau0, clip_below=dt, savepath=SAVE_PATH)
net_settings = dict(N=N, p=p, N_input=us_trn.shape[1], savepath=SAVE_PATH, savename='net')
# initialize net
A_adj, w_in, w_rec = hu.generate_net(**net_settings)
# sampling taus
tau_hom = hu.sample_taus(std = 0*tau0, savename='tau_hom', **tau_settings)
tau_het = hu.sample_taus(std = h*tau0, savename='tau_get', **tau_settings)
# integrating homog
int_settings = dict(A_adj=A_adj, w_in=w_in, w_rec=w_rec, af=af, dt=dt)
int_settings.update(INTEGRATION_SETTINGS[dataset_name])
vs_trn_hom = hu.integrate(u=us_trn, taus=tau_hom, **int_settings)
vs_tst_hom = hu.integrate(u=us_tst, taus=tau_hom, **int_settings)
xs_trn_hom = hu.activation_func(vs_trn_hom, af)
xs_tst_hom = hu.activation_func(vs_tst_hom, af)
# integrating heterog
vs_trn_het = hu.integrate(u=us_trn, taus=tau_het, **int_settings)
vs_tst_het = hu.integrate(u=us_tst, taus=tau_het, **int_settings)
xs_trn_het = hu.activation_func(vs_trn_het, af)
xs_tst_het = hu.activation_func(vs_tst_het, af)
if plot:
# classes distribution
hu.plot_dataset_hist(ys_trn, savepath=SAVE_PATH, savename='trn')
hu.plot_dataset_hist(ys_tst, savepath=SAVE_PATH, savename='tst')
# uvx
hu.plot_uvx_stream(us_trn, vs_trn_hom, xs_trn_hom, lengths=lens_trn,
savename='uvx_homog', savepath=SAVE_PATH,)
hu.plot_uvx_stream(us_trn, vs_trn_het, xs_trn_het, lengths=lens_trn,
savename='uvx_heter', savepath=SAVE_PATH,)
# manifolds
hu.plot_manifold(us_trn, vs_trn_hom, xs_trn_hom, ys_trn, lens_trn,
savename='manifold_homog', savepath=SAVE_PATH,)
hu.plot_manifold(us_trn, vs_trn_het, xs_trn_het, ys_trn, lens_trn,
savename='manifold_heter', savepath=SAVE_PATH,)
# extracting training slice
if num_snapshot>0:
snap_config = dict(nsnapshot=num_snapshot, flatten=True)
_, Xs_trn_hom = hu.get_snapshot(lens_trn, arr=xs_trn_hom, **snap_config)
_, Xs_trn_het = hu.get_snapshot(lens_trn, arr=xs_trn_het, **snap_config)
_, Xs_tst_hom = hu.get_snapshot(lens_tst, arr=xs_tst_hom, **snap_config)
_, Xs_tst_het = hu.get_snapshot(lens_tst, arr=xs_tst_het, **snap_config)
else:
Xs_trn_hom = xs_trn_hom.reshape(len(lens_trn), -1)
Xs_trn_het = xs_trn_het.reshape(len(lens_trn), -1)
Xs_tst_hom = xs_tst_hom.reshape(len(lens_tst), -1)
Xs_tst_het = xs_tst_het.reshape(len(lens_tst), -1)
# classifying
clf_hom = hu.classifier()
clf_het = hu.classifier()
clf_hom.fit(Xs_trn_hom, ys_trn)
clf_het.fit(Xs_trn_het, ys_trn)
# report
print(f'Homog : {clf_hom.score(Xs_trn_hom, ys_trn):.3f} | {clf_hom.score(Xs_tst_hom, ys_tst):.3f}')
print(f'Hetero: {clf_het.score(Xs_trn_het, ys_trn):.3f} | {clf_het.score(Xs_tst_het, ys_tst):.3f}')
[docs]
def run_multinet(
Nnet = 20, # number of MC realization
h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
N=250, p=.1, af = 'tanh', dt=1, # nework parameters
dataset_name='speechcommands', Nbatch=50, top_n='all', # task parameters
unit_std=False, center=False, # input normalization
orthgonalize=False, zero_phase=True, whiten=False, # input orthogonaliziation
normed_trial=True, filt=False, # filtering before tau finding
first_unstream=True, flattening_method='flat', # to stream or not
pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
golength=10, nsnapshot = 10, # pooling args
seed=1234, n_jobs=-1, cv=10, Creg=1,
prefix='', # simulation name
plot_uvx=True, plot_man=True,
**kwargs
):
if prefix == '':
mult = 'M' if h_is_multiplicative else ''
prefix = f'{dataset_name}_N{N}_MC{Nnet}_h{h}{mult}'
prefix+= f'_c{int(center)}_std{int(unit_std)}'
prefix+= f'_o{int(orthgonalize)}_z{int(zero_phase)}_w{int(whiten)}'
print(prefix)
savepath = osjoin(SAVE_PATH, prefix)
Path(savepath).mkdir(parents=True, exist_ok=True)
trn, tst = get_train_test(dataset_name=dataset_name, Nbatch=Nbatch, top_n=top_n,
unit_std=unit_std, center=center, normed_trial=normed_trial,
orthgonalize=orthgonalize, zero_phase=zero_phase)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
Nclasses = len(set(ys_trn).union(set(ys_tst)))
reset_dim = INTEGRATION_SETTINGS[dataset_name].get('reset_dim', 0)
u_reset = INTEGRATION_SETTINGS[dataset_name].get('u_reset', None)
# h = hu.set_het_level(value=h, lens = lens_trn)
# extract tau
tau0, fmaxs = hu.extract_base_tau(us_trn, dt=dt,
sample_lengths=lens_trn,
filt=filt, first_unstream=first_unstream,
flattening_method=flattening_method, center=False,
savepath=savepath, savename='base_extraction'
)
# needed to make sample_integrate_classify_save utility happy :)
tau0s = [tau0]
num_tau0 = 1
pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
memmap_temp = osjoin(savepath, 'joblib_memmap')
Path(memmap_temp).mkdir(parents=True, exist_ok=True)
Parallel(
n_jobs=n_jobs,
backend='loky',
batch_size=1,
pre_dispatch='n_jobs',
max_nbytes=0,
mmap_mode='r',
temp_folder=memmap_temp,
)(
delayed(hu.sample_integrate_classify_save)(
dt=dt, tau0=tau0,
# std=(h*tau0 if h_is_multiplicative else h*dt), dt=dt, N=N, p=p,
h_value=h, h_is_multip = h_is_multiplicative,
us_trn=us_trn, us_tst=us_tst,
lens_trn=lens_trn, lens_tst=lens_tst,
ys_trn=ys_trn, ys_tst=ys_tst,
pooling_funcs=pooling_funcs,
pooling_kwargs=deepcopy(pooling_kwargs),
af=af, cv=cv, Creg=Creg,
prefix=prefix, savepath=savepath,
plot_uvx=plot_uvx, plot_man=plot_man,
net_idx=net_idx, tau0_idx=tau0_idx,
reset_dim=reset_dim, u_reset=u_reset)
for tau0_idx, tau0 in enumerate(tau0s)
for net_idx in range(Nnet)
)
# summarize
trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
savepath=savepath, Ntau0=num_tau0,
Nnet=Nnet, Npool=len(pooling_funcs))
hu.plot_score_summary(
score_trn_hom = trn_homs,
score_trn_het = trn_hets,
score_tst_hom = tst_homs,
score_tst_het = tst_hets,
savepath=savepath, savename=prefix,
chance_score=1./Nclasses,
pooling_funcs=pooling_funcs,
)
[docs]
def run_spaced_tau0_multinet(
Nnet = 20, # number of MC realization
h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
N=250, p=.1, af = 'tanh', J=1, Ju=1, dt=1, # nework parameters
dataset_name='speechcommands', Nbatch=50, top_n=10, # task parameters
unit_std=False, center=False, # input normalization
orthgonalize=False, zero_phase=True, whiten=False, # input orthogonaliziation
normed_trial = True,
filt=False, # filtering before tau finding
first_unstream=True, flattening_method='flat', # to stream or not
num_tau0=10, # number of distinct E[taus] to use
pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
golength=10, nsnapshot = 10, # pooling args
seed=1234, n_jobs=-1, cv=10, Creg=1, clf_name='logreg',
prefix='', # simulation name
plot_uvx=True, plot_man=False,
**kwargs
):
if prefix == '':
mult = 'M' if h_is_multiplicative else ''
prefix = f'SPACED_{dataset_name}_N{N}_MC{Nnet}_h{h}{mult}'
print(prefix)
savepath = osjoin(SAVE_PATH, prefix)
Path(savepath).mkdir(parents=True, exist_ok=True)
trn, tst = get_train_test(dataset_name=dataset_name, Nbatch=Nbatch, top_n=top_n,
unit_std=unit_std, center=center, normed_trial=normed_trial,
orthgonalize=orthgonalize, zero_phase=zero_phase)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
Nclasses = len(set(ys_trn).union(set(ys_tst)))
reset_dim = INTEGRATION_SETTINGS[dataset_name].get('reset_dim', 0)
u_reset = INTEGRATION_SETTINGS[dataset_name].get('u_reset', None)
# h = hu.set_het_level(value=h, lens = lens_trn)
# mean taus values
spaced_tau = hu.get_spaced_hom_taus(lens_trn, dt=dt, num=num_tau0)
pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
memmap_temp = osjoin(savepath, 'joblib_memmap')
Path(memmap_temp).mkdir(parents=True, exist_ok=True)
Parallel(
n_jobs=n_jobs,
backend='loky',
batch_size=1,
pre_dispatch='n_jobs',
max_nbytes=0,
mmap_mode='r',
temp_folder=memmap_temp,
)(
delayed(hu.sample_integrate_classify_save)(
dt=dt, tau0=tau0,
# std=(h*tau0 if h_is_multiplicative else h*dt), dt=dt, N=N, p=p,
h_value=h, h_is_multip = h_is_multiplicative,
N=N, p=p, J=J, Ju=Ju,
us_trn=us_trn, us_tst=us_tst,
lens_trn=lens_trn, lens_tst=lens_tst,
ys_trn=ys_trn, ys_tst=ys_tst,
pooling_funcs=pooling_funcs,
pooling_kwargs=deepcopy(pooling_kwargs),
af=af, cv=cv, Creg=Creg, clf_name=clf_name,
prefix=prefix, savepath=savepath,
plot_uvx=plot_uvx, plot_man=plot_man,
net_idx=net_idx, tau0_idx=tau0_idx,
reset_dim=reset_dim, u_reset=u_reset)
for tau0_idx, tau0 in enumerate(spaced_tau)
for net_idx in range(Nnet)
)
# summarize
trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
savepath=savepath, Ntau0=num_tau0,
Nnet=Nnet, Npool=len(pooling_funcs))
hu.plot_score_summary_of_spaced_hom(trn_homs, trn_hets, tst_homs, tst_hets,
spaced_taus_hom = spaced_tau,
savepath=savepath, savename=prefix,
chance_score=1./Nclasses,
pooling_funcs=pooling_funcs, Cregs=Creg,
)
het_taus = hu.collect_taus(savepath=savepath, Ntau0=len(spaced_tau), Nnet=Nnet)
hu.plot_taus_summary(het_taus=het_taus, hom_taus=spaced_tau, savepath=savepath)
if __name__ == '__main__':
print(torch.__version__)
print(torchaudio.__version__)
print(f'device is: {device}')
print(f'data path is {datapath}')
parser = argparse.ArgumentParser()
parser.add_argument('dataset_name', type=str, default=None)
parser.add_argument('Nbatch', type=int, default=100)
parser.add_argument('is_spaced', type=bool, default=True)
parser.add_argument('n_jobs', type=int, default=10)
args = parser.parse_args()
if args.is_spaced:
run_spaced_tau0_multinet(Nbatch=args.Nbatch, n_jobs=args.n_jobs,
dataset_name=args.dataset_name)
else:
run_multinet(Nbatch=args.Nbatch, n_jobs=args.n_jobs,
dataset_name=args.dataset_name)
#run_multinet(Nbatch=50, dataset_name='mnist', n_jobs=6)
# run_spaced_tau0_multinet(Nbatch=100,)
# run_multiple_net()
#run_spaced_hom(N=250, dataset_name='mnist', Nbatch=100, h='paper', h_is_multiplicative=True)
# run_spaced_hom(N=250, dataset_name='speechcommands', Nbatch=100, h='paper', h_is_multiplicative=True)
# run_spaced_hom(N=100, dataset_name='speechcommands', Nbatch=100, h='paper', h_is_multiplicative=True)
# run_spaced_hom(N=50, dataset_name='speechcommands', Nbatch=100, h='paper', h_is_multiplicative=True)