def on_hidacts_raw(hid_acts: torch.tensor):
if hid_acts.shape[1] >= num_pcs:
pc_vals, pc_vecs = pca.get_hidden_pcs(hid_acts, num_pcs)
projected = pca.project_to_pcs(hid_acts, pc_vecs, out=None)
snapshots.append(PCTrajectoryGenSnapshot(pc_vecs, pc_vals, projected, out))
elif hid_acts.shape[1] > 1:
pc_vals, pc_vecs = pca.get_hidden_pcs(hid_acts, hid_acts.shape[1])
projected = pca.project_to_pcs(hid_acts, pc_vecs, out=None)
pc_vecs_app = torch.zeros((num_pcs, hid_acts.shape[1]), dtype=pc_vecs.dtype)
pc_vecs_app[:pc_vecs.shape[0]] = pc_vecs
pc_vals_app = torch.zeros((num_pcs,), dtype=pc_vals.dtype)
pc_vals_app[:pc_vals.shape[0]] = pc_vals
projected_app = torch.zeros((hid_acts.shape[0], num_pcs), dtype=projected.dtype)
projected_app[:projected.shape[0]] = projected
snapshots.append(PCTrajectoryGenSnapshot(pc_vecs_app, pc_vals_app, projected_app, out))
else:
pc_vecs = torch.zeros((num_pcs, hid_acts.shape[1]), dtype=hid_acts.dtype)
pc_vals = torch.zeros((num_pcs,), dtype=hid_acts.dtype)
pc_vals[0] = 1
pc_vecs[0, 0] = 1
projected = torch.zeros((hid_acts.shape[0], num_pcs), dtype=hid_acts.dtype)
projected[:, 0] = hid_acts.squeeze()
snapshots.append(PCTrajectoryGenSnapshot(pc_vecs, pc_vals, projected, out))
def measure_pr(hidden_acts: torch.tensor) -> float:
"""Measures the participation ratio of the specified hidden acts, plotting
it as necessary and saving it to the zip file in the given directory.
Args:
hidden_acts (torch.tensor): a [batch_size x layer_size] tensor of activations,
where the first index tells us which point and the second index tells us
which neuron. float-like
"""
if not torch.is_tensor(hidden_acts):
raise ValueError(f'expected hidden_acts is tensor, got {hidden_acts} (type={type(hidden_acts)})')
if len(hidden_acts.shape) != 2:
raise ValueError(f'expected hidden_acts.shape is (batch_size, layer_size), got {hidden_acts.shape}')
if hidden_acts.dtype not in (torch.float, torch.double):
raise ValueError(f'expected hidden_acts is float-like, but dtype is {hidden_acts.dtype}')
eigs = pca.get_hidden_pcs(hidden_acts, None, False)
result = torch.pow(torch.sum(eigs), 2) / torch.sum(torch.pow(eigs, 2))
result = result.item()
if not isinstance(result, float):
raise ValueError(f'expected participation_ratio is float, got {result} (type={type(result)})')
if result < 1 - 1e-8:
raise ValueError(f'PR should always be >= 1 (got {result})')
if result > eigs.shape[0] + 1e-8:
raise ValueError(f'PR should always be <= M where M = eigs.shape[0] = {eigs.shape[0]} (got {result})')
return result
def digest_find_and_plot_traj(
sample_points: np.ndarray,
sample_labels: np.ndarray, # pylint: disable=unused-argument
*all_hid_acts: typing.Tuple[np.ndarray],
savepath: str = None,
**kwargs):
"""Digestor friendly way to find the pc trajectory and then plot it at the given
savepath for the given sample points, labels, and hidden activations"""
sample_points = torch.from_numpy(sample_points)
sample_labels = torch.from_numpy(sample_labels)
hacts_cp = []
for hact in all_hid_acts:
hacts_cp.append(torch.from_numpy(hact))
all_hid_acts = hacts_cp
snapshots = []
for hid_acts in all_hid_acts:
pc_vals, pc_vecs = get_hidden_pcs(hid_acts, 2)
projected = project_to_pcs(hid_acts, pc_vecs, out=None)
snapshots.append(
PCTrajectoryFFSnapshot(pc_vecs, pc_vals, projected, sample_labels))
traj = PCTrajectoryFF(snapshots)
plot_trajectory(traj, savepath, False, **kwargs)
def on_hidacts(acts_info: FFHiddenActivations):
hid_acts = acts_info.hidden_acts.detach()
pc_vals, pc_vecs = get_hidden_pcs(hid_acts, num_pcs)
projected = project_to_pcs(hid_acts, pc_vecs, out=None)
snapshots.append(
PCTrajectoryFFSnapshot(pc_vecs, pc_vals, projected, sample_labels))
def _update_match(self):
pc_vals, pc_vecs = pca.get_hidden_pcs(self.hidacts_torch, 3)
projected = pca.project_to_pcs(self.hidacts_torch, pc_vecs, out=None)
snap = pca_ff.PCTrajectoryFFSnapshot(pc_vecs, pc_vals, projected,
self.sample_labels_torch)
self.match_info.match(snap)
match_info = pca_ff.PCTrajectoryFFSnapshotMatchInfo.create(snap)
self.match_mean_comps_torch[:] = match_info.mean_comps
def _get_snapshot(self):
pc_vals, pc_vecs = pca.get_hidden_pcs(self.hidden_acts_torch, 3)
projected = pca.project_to_pcs(self.hidden_acts_torch,
pc_vecs,
out=None)
snap = pca_ff.PCTrajectoryFFSnapshot(pc_vecs, pc_vals, projected,
self.sample_labels_torch)
self.match_info.match(snap)
self.snapshot = snap
def to_trajectory(sample_labels: torch.tensor,
all_hid_acts: typing.List[torch.tensor],
num_pcs: int) -> PCTrajectoryFF:
"""Converts the specified hidden activations to a feedforward trajectory
Args:
sample_labels (torch.tensor): the labels for the points sent through the model
all_hid_acts (typing.List[torch.tensor]): the hidden activations of the network
num_pcs (int): the number of hidden pcs to find
Returns:
PCTrajectoryFF: the projectory formed by the specified hidden activations
"""
if not torch.is_tensor(sample_labels):
raise ValueError(
f'expected sample_labels is tensor, got {sample_labels} (type={type(sample_labels)})'
)
if sample_labels.dtype not in (torch.uint8, torch.int, torch.long):
raise ValueError(
f'expected sample_labels is int-like but has dtype {sample_labels.dtype}'
)
if len(sample_labels.shape) != 1:
raise ValueError(
f'expected sample_labels has shape [batch_size] but is {sample_labels.shape}'
)
if not isinstance(all_hid_acts, (tuple, list)):
raise ValueError(
f'expected all_hid_acts is tuple or list, got {all_hid_acts} (type={type(all_hid_acts)})'
)
snapshots = []
for idx, hid_acts in enumerate(all_hid_acts):
if not torch.is_tensor(hid_acts):
raise ValueError(
f'expected all_hid_acts[{idx}] is tensor, got {hid_acts} (type={type(hid_acts)})'
)
if hid_acts.dtype not in (torch.float, torch.double):
raise ValueError(
f'expected all_hid_acts[{idx}] is float-like but has dtype {hid_acts.dtype}'
)
if len(hid_acts.shape
) != 2 or hid_acts.shape[0] != sample_labels.shape[0]:
raise ValueError(
f'expected all_hid_acts[{idx}].shape = [batch_size={sample_labels.shape[0]}, layer_size] but got {hid_acts.shape}'
)
pc_vals, pc_vecs = get_hidden_pcs(hid_acts, num_pcs)
projected = project_to_pcs(hid_acts, pc_vecs, out=None)
snapshots.append(
PCTrajectoryFFSnapshot(pc_vecs, pc_vals, projected, sample_labels))
return PCTrajectoryFF(snapshots)
def to_trajectory(sample_labels: torch.tensor, all_hid_acts: typing.List[torch.tensor],
num_pcs: int) -> PCTrajectoryGen:
"""Converts the given labels and hidden activations to a trajectory with the specified
number of principal components.
Args:
sample_labels (torch.tensor): The labels / targets for the corresponding samples
size: [num_samples, ...]
dtype: any
first index is which sample
all_hid_acts (list[torch.tensor]): the hidden activations through time or layers
size: [num_layers]
each element:
size: [num_samples, layer_size]
dtype: float-like
the first index is which sample
the second index is which feature
Returns:
The trajectory found by calculating the top num_pcs pcs and then projecting the samples
onto those pcs. For plotting it may be helpful to use a PCTrajectoryFFMatchInfo which works
for the Gen version as well (since it is agnostic to labels)
"""
tus.check(sample_labels=(sample_labels, torch.tensor),
all_hid_acts=(all_hid_acts, (list, tuple)),
num_pcs=(num_pcs, int))
tus.check_listlike(all_hid_acts=(all_hid_acts, torch.Tensor))
snapshots = []
for idx, hid_acts in enumerate(all_hid_acts):
torch.check_tensors(
**{f'hid_acts[{idx}]': (
hid_acts,
(
('num_samples', sample_labels.shape[0]),
('layer_size', None)
),
(torch.float, torch.double)
)}
)
pc_vals, pc_vecs = pca.get_hidden_pcs(hid_acts, num_pcs)
projected = pca.project_to_pcs(hid_acts, pc_vecs, out=None)
snapshots.append(PCTrajectoryGenSnapshot(pc_vecs, pc_vals, projected, sample_labels))
return PCTrajectoryGen(snapshots)