""" ate.py # ------------------------------------------- Intervene at fixed timestep T on the trained LSTM model, clamp a unit to a value and get new rating. Get change in rating after intervention, loop over different clamp values, and then for all units, and then for all frames. Specify stimuli population across which the ratings are meaned. Produces: 1. Temporal causal heatmap (absolute ATE ranges for all units all frames) 2. ATE curves plot (all 64 units for one frame, one unit highlighted) """ import torch import torchaudio import torch.nn as nn import numpy as np import pandas as pd import matplotlib.pyplot as plt import matplotlib.cm as cm import matplotlib.colors as mcolours from i_dataset import StimuliSounds from ii_structure import SimpleLSTM # Configuration (same as trainiing) # ------------------------------------------- CSV_FILE = "../../dataset.csv" AUDIO_DIR = "../../sound/" SAMPLES = 6394 INPUT_SIZE = 128 HIDDEN_SIZE = 64 LEARNING_RATE = 0.01 MODEL_PATH = "model.pt" mel_transform = nn.Sequential( torchaudio.transforms.MelSpectrogram( sample_rate=16000, n_fft=1024, win_length=320, hop_length=44, # 2.75 ms per frame window_fn=torch.hamming_window, n_mels=128 ), torchaudio.transforms.AmplitudeToDB() ) # Frame duration and intervention frame FRAME_DURATION_MS = 44 / 16000 * 1000 # 2.75 ms INTERVENTION_FRAME = 90 # 55 = 151.25, 73 = 200.75 ms, 90 = 247.5, 91 = 250.25 # Load model and dataset # ------------------------------------------- model = SimpleLSTM(input_size=INPUT_SIZE, hidden_size=HIDDEN_SIZE, learning_rate=LEARNING_RATE) model.load_state_dict(torch.load(MODEL_PATH)) dataset = StimuliSounds(CSV_FILE, AUDIO_DIR, SAMPLES, mel_transform) csv_data = pd.read_csv(CSV_FILE) # File configurations # ------------------------------------------- unique_files = csv_data.iloc[:, 0].unique().tolist() AsEfiles = [f for f in unique_files if 'asm' in f] AmEfiles = [f for f in unique_files if 'amr' in f] Cɔtfiles = [f for f in unique_files if '1_' in f] Cruːfiles = [f for f in unique_files if '2_' in f] AsE1files = [f for f in unique_files if 'asm' in f and '1_' in f] AsE2files = [f for f in unique_files if 'asm' in f and '2_' in f] AmE1files = [f for f in unique_files if 'amr' in f and '1_' in f] AmE2files = [f for f in unique_files if 'amr' in f and '2_' in f] # Choose stimuli population (remember to change plot labels) stimulifiles = unique_files # Intervention and ATE functions # ------------------------------------------- def intervention(model, mel, T, unit_idx, clamp_value): """ Run LSTM on mel up to frame T, clamp hidden unit unit_idx to a clamp_value, then continue LSTM for rest of frames. Returns final predicted rating. """ model.eval() # tells pytorch to switch to model evaluation mode with torch.no_grad(): # initial hidden states (zeros) h = torch.zeros(1, 1, model.lstm.hidden_size) c = torch.zeros(1, 1, model.lstm.hidden_size) # (1, 1, 64) # process frames from 0 up to and including T beforeframes = mel[:, :T+1, :] # (1, T+1, n_mels) = (batch, frames, mels) _, (h, c) = model.lstm(beforeframes, (h, c)) # h,c: (1,1,hidden) # intervene h_clamped = h.clone() h_clamped[0, 0, unit_idx] = clamp_value # process rest of frames (T+1 onward) afterframes = mel[:, T+1:, :] # (1, remaining, n_mels) if afterframes.shape[1] > 0: _, (h_clamped, c) = model.lstm(afterframes, (h_clamped, c)) # Final rating from last hidden state rating = model.fc(h_clamped[0, 0]).item() # indexes into first layer, first batch giving (64,) # model.fc makes (64,) > (1,) # .item() returns the signle value as std python float return rating def compute_ate_curve(model, dataset, file_list, T, unit_idx, clamp_values): """ Get mean of ratings across stimuli on trained LSTM without intervention, and then the same with intervention. Calculate ate values for each clamped value. """ # Baseline: natural rating (no intervention) natural_ratings = [] for fname in file_list: idx = csv_data[csv_data.iloc[:, 0] == fname].index[0] # get row num mel, _ = dataset[idx] # pass it to StimuliSounds mel = mel.unsqueeze(0) natural_ratings.append(model.forward(mel).item()) # list of all baseline ratings baseline = np.mean(natural_ratings) ate_values = [] for val in clamp_values: clamped_ratings = [] for fname in file_list: idx = csv_data[csv_data.iloc[:, 0] == fname].index[0] mel, _ = dataset[idx] mel = mel.unsqueeze(0) rating = intervention(model, mel, T, unit_idx, val) clamped_ratings.append(rating) ate = np.mean(clamped_ratings) - baseline ate_values.append(ate) # avg rating of the population (= all stimuli or a section) << baseline, run model normally # run model with intervention for a given unit at a give frame T # >> get clamped ratings for stimuli >> get avg >> get ate # (loops for units and frames handled separately) # >> get ate for different clamp values for ate curve >> ate_values return ate_values # Temporal causal heatmap (ATE range at every 5 frames) # ------------------------------------------- # Choose frame step and clamp range time_frames = np.arange(1, 145, 5) # starts at 1, step 5, stops before 145 clamp_values = np.linspace(-10, 10, 11) # or np.linspace(-2, 11, 11) causal_map = np.zeros((HIDDEN_SIZE, len(time_frames))) print("Computing temporal causal heatmap...") # takes time for t_idx, T in enumerate(time_frames): # loop over all T in time_frames print(f" Frame {T}/{145}") for unit in range(HIDDEN_SIZE): # ATE curve (list of floats all units) ate_vals = compute_ate_curve(model, dataset, stimulifiles, T, unit, clamp_values) # Summarise as absolute ATE range ate_range = abs(max(ate_vals) - min(ate_vals)) causal_map[unit, t_idx] = ate_range # hold the absolute ATE range for a unit at time_frames[t_idx] # Plot the heatmap plt.figure(figsize=(12, 8)) plt.imshow(causal_map, aspect='auto', origin='lower', extent = [time_frames[0] * FRAME_DURATION_MS, time_frames[-1] * FRAME_DURATION_MS, 0, HIDDEN_SIZE - 1], cmap='viridis') plt.yticks(np.arange(0, HIDDEN_SIZE, 5)) # yaxis ticks for every 5 units plt.colorbar(label='|ATE range|') plt.xlabel('Intervention time (ms)') plt.ylabel('Hidden unit index') plt.title('Temporal causal influence of each hidden unit (ATE range) for all stimuli') plt.tight_layout() plt.savefig("temporal-ate-All.png") plt.show() # ATE scanning over all units # ------------------------------------------- clamp_values = np.linspace(-10.0, 10.0, 31) # 31 points print("Computing ATE curves for all units...") # takes some time all_ate_curves = [] for unit in range(HIDDEN_SIZE): ate = compute_ate_curve(model, dataset, stimulifiles, INTERVENTION_FRAME, unit, clamp_values) all_ate_curves.append(ate) # Optional: print progress if (unit+1) % 8 == 0: print(f" {unit+1}/{HIDDEN_SIZE}") highlighted_unit = 38 cmap = cm.viridis norm = mcolours.Normalize(vmin=0, vmax=63) fig, (ax_left, ax_right) = plt.subplots(1, 2, figsize=(16, 6), sharey=True) # Left facet: all units coloured for unit, ate in enumerate(all_ate_curves): colour = cmap(norm(unit)) ax_left.plot(clamp_values, ate, color=colour, lw=1) sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) sm.set_array([]) fig.colorbar(sm, ax=ax_left, label='Unit index') ax_left.axhline(0, color='black', linestyle='--') ax_left.set_xlabel('Clamped value') ax_left.set_ylabel('ATE (change in rating from baseline)') ax_left.set_title('All units') ax_right.legend() # Right facet: greyed units + highlighted unit for unit, ate in enumerate(all_ate_curves): if unit == highlighted_unit: continue ax_right.plot(clamp_values, ate, color='lightgrey', alpha=0.5, lw=1) # Highlighted unit ate_highlight = all_ate_curves[highlighted_unit] ax_right.plot(clamp_values, ate_highlight, color='red', lw=2.5, label=f'Unit {highlighted_unit}') ax_right.axhline(0, color='black', linestyle='--') ax_right.set_xlabel('Clamped value') ax_right.set_title(f'Unit {highlighted_unit}') # Overall title fig.suptitle(f'Intervention at frame {INTERVENTION_FRAME} ' f'(~{INTERVENTION_FRAME * FRAME_DURATION_MS:.0f} ms)', fontsize=14) plt.tight_layout() plt.savefig("populationATEcurve-90.png", dpi=300, bbox_inches='tight') plt.show()