def create_loaders(datasets,
args,
split_test=True,
test_size=1,
context_filter=[]):
dsets_train = []
dsets_test = []
for i, dpath in enumerate(datasets):
dset = load_rb(dpath)
# trim and set name of each dataset
dname = str(i) + '_' + ':'.join(dpath.split('/')[-1].split('.')[:-1])
if split_test:
cutoff = round(.9 * len(dset))
dsets_train.append([dname, dset[:cutoff]])
dsets_test.append([dname, dset[cutoff:]])
else:
dsets_test.append([dname, dset])
# creating datasets
test_set = TrialDataset(dsets_test, args)
if split_test:
train_set = TrialDataset(dsets_train, args)
# TODO: make all this code better
if args.sequential:
# helper function for sequential loaders
def create_subset_loaders(dset, batch_size, drop_last):
loaders = []
max_idxs = dset.max_idxs
for i in range(len(datasets)):
if i == 0:
subset = Subset(dset, range(max_idxs[0]))
else:
subset = Subset(dset, range(max_idxs[i - 1], max_idxs[i]))
loader = DataLoader(subset,
batch_size=batch_size,
shuffle=True,
collate_fn=collater,
drop_last=drop_last)
loaders.append(loader)
return loaders
# create the loaders themselves
test_loaders = create_subset_loaders(test_set, test_size, False)
if split_test:
train_loaders = create_subset_loaders(train_set, args.batch_size,
True)
return (train_set, train_loaders), (test_set, test_loaders)
return (test_set, test_loaders)
# filter out some contexts
elif len(context_filter) != 0:
def create_context_loaders(dset, batch_size, drop_last):
max_idxs = dset.max_idxs
c_range = []
for i in range(len(datasets)):
if i in context_filter:
continue
if i == 0:
c_range += list(range(max_idxs[0]))
else:
c_range += list(range(max_idxs[i - 1], max_idxs[i]))
subset = Subset(dset, c_range)
loader = DataLoader(subset,
batch_size=batch_size,
shuffle=True,
collate_fn=collater,
drop_last=drop_last)
return loader
# create the loaders themselves
test_loaders = create_context_loaders(test_set, test_size, False)
if split_test:
train_loaders = create_context_loaders(train_set, args.batch_size,
True)
return (train_set, train_loaders), (test_set, test_loaders)
return (test_set, test_loaders)
else:
# otherwise it's quite simple, create a single dataset and loader
test_loader = DataLoader(test_set,
batch_size=test_size,
shuffle=True,
collate_fn=collater,
drop_last=False)
if split_test:
train_loader = DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
collate_fn=collater,
drop_last=True)
return (train_set, train_loader), (test_set, test_loader)
return (test_set, test_loader)
def test_model(net, config, n_tests=0):
dset = load_rb(config.dataset)
dset_idx = range(len(dset))
p_fn = get_potential(config)
criterion = nn.MSELoss()
if n_tests != 0:
dset_idx = sorted(random.sample(range(len(dset)), n_tests))
dset = [dset[i] for i in dset_idx]
is_goals_task = config.dset_type == 'goals'
x, y = get_x_y(dset, config.dset_type)
with torch.no_grad():
net.reset()
losses = []
outs = []
if is_goals_task:
ins = []
n_pts = x.shape[1]
cur_idx = torch.zeros(x.shape[0], dtype=torch.long)
for j in range(config.goals_timesteps):
net_in = x[torch.arange(x.shape[0]),
cur_idx, :] #.reshape(-1, net.args.L)
ins.append(net_in)
#pdb.set_trace()
net_out, extras = net(net_in, extras=True)
# net_target = net_in.reshape(-1, net.args.Z)
trial_losses = []
# dones = torch.zeros(x.shape[0], dtype=torch.long)
for k in range(len(net_out)):
# need to add the dimension back in so the goals loss fn works
net_out_k = net_out[k].unsqueeze(0)
x_k = x[k].unsqueeze(0)
# need to adjust this because we're separating losses from each other
step_loss, cur_idx[k] = goals_loss(
net_out_k,
x_k,
cur_idx[k],
threshold=config.goals_threshold)
trial_losses.append(step_loss)
# dones[k] = done.item()
# cur_idx = update_seq_indices(x, cur_idx, dones)
losses.append(np.array(trial_losses))
outs.append(net_out)
print(cur_idx)
goals = x
ins = torch.stack(ins, dim=1).squeeze()
else:
for j in range(x.shape[1]):
# run the step
net_in = x[:, j].reshape(-1, net.args.L)
net_out = net(net_in)
outs.append(net_out)
net_target = y[:, j].reshape(-1, net.args.Z)
trial_losses = []
for k in range(len(dset)):
step_loss = criterion(net_out[k], net_target[k])
trial_losses.append(step_loss)
losses.append(np.array(trial_losses))
ins = x
goals = x
losses = np.sum(losses, axis=0)
z = torch.stack(outs, dim=1).squeeze()
data = list(zip(dset_idx, ins, goals, z, losses))
final_loss = np.mean(losses, axis=0)
return data, final_loss
for j, dset in enumerate(dsets):
subset = dt[dt.dset == dset]
for iterr in range(len(subset)):
job_id = subset.iloc[iterr].slurm_id
model_folder = os.path.join('..', 'logs', run_id, str(job_id))
model_path = os.path.join(model_folder, 'model_best.pth')
config = get_config(model_path, ctype='model', to_bunch=True)
config.m_noise = 0
config.dataset = dset_map[config.dataset]
net = load_model_path(model_path, config=config)
data, loss = test_model(net, config, n_tests=200, dset_base='../')
dset = load_rb(os.path.join('..', config.dataset))
distr = {}
for k in range(len(data)):
dset_idx, x, _, z, _ = data[k]
r, s, g = dset[dset_idx][2]
t_first = torch.nonzero(z >= 1)
if len(t_first) > 0:
t_first = t_first[0, 0]
else:
t_first = len(x)
val = np.array(t_first - s - 5)
del config_args.config
args = update_args(args, config_args)
else:
# add task-specific arguments. shouldn't need to do this if loading from config file
task_args = get_task_args(args)
args = update_args(args, task_args)
args.argv = ' '.join(sys.argv)
if args.mode == 'create':
# create and save a dataset
dset, config = create_dataset(args)
save_dataset(dset, args.name, config=config)
elif args.mode == 'load':
# visualize a dataset
dset = load_rb(args.name)
t_type = type(dset[0])
xr = np.arange(dset[0].t_len)
samples = random.sample(dset, 12)
fig, ax = plt.subplots(3, 4, sharex=True, sharey=True, figsize=(10, 6))
for i, ax in enumerate(fig.axes):
ax.axvline(x=0, color='dimgray', alpha=1)
ax.axhline(y=0, color='dimgray', alpha=1)
ax.grid(True, which='major', lw=1, color='lightgray', alpha=0.4)
ax.tick_params(axis='both', color='white')
#ax.set_title(sample[i][2])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_visible(False)
import matplotlib.pyplot as plt
import argparse
import pdb
import os
import sys
sys.path.append('../')
from utils import Bunch, load_rb
from train import Trainer, parse_args, adjust_args
from network import HypothesisNetwork
b = Bunch()
b.dataset = '../datasets/temp.pkl'
b.out_act = 'none'
b.L = 2
b.Z = 2
dset = load_rb(b.dataset)
net = HypothesisNetwork(b)
seq1 = dset[0]
t1 = seq1[0]
t1 = torch.Tensor(t1)
net(t1)
with open(args.model, 'rb') as f:
model = torch.load(f)
if args.noise != 0:
J = model['W_f.weight']
v = J.std()
shp = J.shape
model['W_f.weight'] += torch.normal(0, v * .5, shp)
J = model['W_ro.weight']
v = J.std()
shp = J.shape
model['W_ro.weight'] += torch.normal(0, v * .5, shp)
net = load_model_path(args.model, params={'dset': args.dataset, 'out_act': args.out_act})
dset = load_rb(args.dataset)
data = test_model(net, dset, n_tests=0)
run_id = '/'.join(args.model.split('/')[-3:-1])
fig, ax = plt.subplots(3,4,sharex=True, sharey=True, figsize=(12,7))
for i, ax in enumerate(fig.axes):
ix, x, y, z, loss = data[i]
xr = np.arange(len(x))
ax.axvline(x=0, color='dimgray', alpha = 1)
ax.axhline(y=0, color='dimgray', alpha = 1)
ax.grid(True, which='major', lw=1, color='lightgray', alpha=0.4)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
def create_dataset(args):
name = args.name
t_type = args.trial_type
n_trials = args.n_trials
t_len = args.trial_len
trial_args = args.trial_args
config = {}
config['argv'] = sys.argv
config['t_type'] = t_type
config['n_trials'] = n_trials
config['t_len'] = t_len
trials = []
if t_type.startswith('copy'):
delay = get_args_val(trial_args, 'delay', 0, int)
config['delay'] = delay
dim = 1
x = np.arange(0, t_len)
ys = []
if t_type == 'copy':
n_freqs = get_args_val(trial_args, 'n_freqs', 15, int)
f_range = get_args_val(trial_args,
'f_range', [2, 30],
float,
n_vals=2)
amp = get_args_val(trial_args, 'amp', 1, float)
start_zero = 'start_nonzero' not in trial_args
config['n_freqs'] = n_freqs
config['f_range'] = f_range
config['amp'] = amp
config['start_zero'] = start_zero
fn = np.sin if start_zero else np.cos
for n in range(n_trials):
y = np.zeros_like(x)
freqs = np.random.uniform(f_range[0], f_range[1], (n_freqs))
amps = np.random.uniform(-amp, amp, (n_freqs))
for i in range(n_freqs):
y = y + amps[i] * fn(1 / freqs[i] * x)
ys.append(y)
elif t_type == 'copy_gp':
interval = get_args_val(trial_args, 'interval', 10, int)
scale = get_args_val(trial_args, 'scale', 1, float)
kernel = RBF(length_scale=5)
config['interval'] = interval
config['scale'] = scale
x_list = x[..., np.newaxis]
x_filter = x_list[::interval]
n_pts = x_filter.squeeze().shape[0]
for n in range(n_trials):
y_filter = np.zeros((n_pts, dim))
y_filter[0] = 0
for i in range(1, n_pts):
if 'smoothing' in trial_args:
y_filter[i] = np.random.multivariate_normal(
y_filter[i - 1] / 2, cov=scale * np.eye(dim))
else:
y_filter[i] = np.random.multivariate_normal(
np.zeros(dim), cov=scale * np.eye(dim))
gp = gpr(kernel=kernel,
normalize_y=False).fit(x_filter, y_filter)
y_prediction, y_std = gp.predict(x_list, return_std=True)
y = y_prediction.reshape(-1)
ys.append(y)
elif t_type == 'copy_motifs':
assert args.motifs is not None
motifs = load_rb(args.motifs)
pause = get_args_val(trial_args, 'pause', 10, int)
amp = get_args_val(trial_args, 'amp', .1, float)
config['pause'] = pause
config['amp'] = amp
for n in range(n_trials):
y = gen_fn_motifs(motifs,
length=t_len,
pause=pause,
amp=amp,
smoothing='cubic')
ys.append(y)
else:
raise Exception
for y in ys:
z = np.zeros_like(y)
if delay == 0:
z = y
else:
z[delay:] = y[:-delay]
trials.append((y, z, delay))
elif t_type == 'integration':
n_freqs = get_args_val(trial_args, 'n_freqs', 15, int)
f_range = get_args_val(trial_args, 'f_range', [3, 30], float, n_vals=2)
amp = get_args_val(trial_args, 'amp', 1, int)
config['n_freqs'] = n_freqs
config['f_range'] = f_range
config['amp'] = amp
for n in range(n_trials):
x = np.arange(0, t_len)
y = np.zeros_like(x).astype(np.float32)
xp = t_len // 2
freqs = np.random.uniform(f_range[0], f_range[1], (n_freqs))
amps = np.random.uniform(-amp, amp, (n_freqs))
for i in range(n_freqs):
y[:xp] = y[:xp] + amps[i] * np.cos(1 / freqs[i] * x[:xp])
y[:xp] = y[:xp] * np.cos(np.pi / 2 * x[:xp] / x[xp])
z_mag = np.sum(y)
trial_range = np.arange(t_len)
z = z_mag / 2 * norm.pdf(trial_range, loc=int(xp * 3 / 2), scale=2)
trials.append((y, z, z_mag))
elif t_type == 'goals':
n_goals = get_args_val(trial_args, 'n_goals', 10, int)
scale = get_args_val(trial_args, 'scale', 5, float)
dim = get_args_val(trial_args, 'dim', 2, int)
config['n_goals'] = n_goals
config['scale'] = scale
config['dim'] = dim
config['goals-type'] = 'normal'
if 'across' in trial_args:
config['goals-type'] = 'across'
for n in range(n_trials):
if 'across' in trial_args:
trials.append(across_goals(n_goals, dim, scale))
else:
for n in range(n_trials):
trial = []
for i in range(n_goals):
trial.append(np.random.normal(loc=0, scale=scale,
size=dim))
trials.append(trial)
return trials, config
def __init__(self, args):
super().__init__()
self.args = args
if self.args.net == 'basic':
self.net = BasicNetwork(self.args)
elif self.args.net == 'state':
self.net = StateNet(self.args)
elif self.args.net == 'hypothesis':
self.net = HypothesisNet(self.args)
# picks which parameters to train and which not to train
self.n_params = {}
self.train_params = []
self.not_train_params = []
logging.info('Training the following parameters:')
for k, v in self.net.named_parameters():
# k is name, v is weight
found = False
# filtering just for the parts that will be trained
for part in self.args.train_parts:
if part in k:
logging.info(f' {k}')
self.n_params[k] = (v.shape, v.numel())
self.train_params.append(v)
found = True
break
if not found:
self.not_train_params.append(k)
logging.info('Not training:')
for k in self.not_train_params:
logging.info(f' {k}')
self.criterion = get_criterion(self.args)
self.optimizer = get_optimizer(self.args, self.train_params)
self.dset = load_rb(self.args.dataset)
self.potential = get_potential(self.args)
# if using separate training and test sets, separate them out
if not self.args.same_test:
np.random.shuffle(self.dset)
cutoff = round(.9 * len(self.dset))
self.train_set = self.dset[:cutoff]
self.test_set = self.dset[cutoff:]
logging.info(
f'Using separate training ({cutoff}) and test ({len(self.dset) - cutoff}) sets.'
)
else:
self.train_set = self.dset
self.test_set = self.dset
self.log_interval = self.args.log_interval
if not self.args.no_log:
self.log = self.args.log
self.run_id = self.args.log.run_id
self.vis_samples = []
self.csv_path = open(
os.path.join(self.log.run_dir, f'losses_{self.run_id}.csv'),
'a')
self.writer = csv.writer(self.csv_path,
delimiter=',',
quotechar='|',
quoting=csv.QUOTE_MINIMAL)
self.writer.writerow(['ix', 'avg_loss'])
self.plot_checkpoint_path = os.path.join(
self.log.run_dir, f'checkpoints_{self.run_id}.pkl')
self.save_model_path = os.path.join(self.log.run_dir,
f'model_{self.run_id}.pth')
args = parser.parse_args()
with open(args.model, 'rb') as f:
m_dict = torch.load(f)
J = m_dict['W_f.weight']
v = J.std()
shp = J.shape
m_dict['W_f.weight'] += torch.normal(0, v * .01, shp)
J = m_dict['W_ro.weight']
v = J.std()
shp = J.shape
m_dict['W_ro.weight'] += torch.normal(0, v * .01, shp)
dset = load_rb(args.dset)
test_data = test_model(m_dict, dset, n_tests=200)
i, xs, ys, zs, losses = list(zip(*test_data))
print(np.mean(losses))
# bunch = Bunch()
# bunch.N = 250
# bunch.D = 250
# bunch.O = 1
# bunch.res_init_type = 'gaussian'
# bunch.res_init_params = {'std': 1.5}
# bunch.reservoir_seed = 0
# net = Network(bunch)
