<jupyter_start><jupyter_text>Deriving Attention Here we are iterating over the whole dataset and calculating attention values used for every auto-regressive prediction<jupyter_code>loader = DataLoader(dataset, shuffle=False, pin_memory=False,
                                  batch_size=1, num_workers=1)

attention_intervals = att_interval(model, loader)
# attention_intervals = np.load(parent_path + "code/transformer_vid3/runs/attention_intervals_n.npy")
# attention_intervals[:, frame_block_size:].shape<jupyter_output><empty_output><jupyter_text>To find some interesting maps, we calculate the standard deviation of each attention step for each interval we are predicting in.<jupyter_code>att_trials_id = attention_intervals[:, frame_block_size:, :frame_block_size]
# att_trials_id = att_trials_id.reshape(-1, att_trials_id.shape[-1])
att_trials_id_std = np.std(att_trials_id, axis=-1)

vid_stack = video_stack.view(3, video_stack.shape[1], video_stack.shape[3], video_stack.shape[4], video_stack.shape[2])
print(vid_stack.shape)<jupyter_output>torch.Size([3, 640, 64, 112, 1])<jupyter_text>We then unravel the flattened attention map back into a 3D (time - length - width) map, and then project that map onto our original image.<jupyter_code>sorted_att_std = np.unravel_index(np.argsort(-att_trials_id_std.ravel()), att_trials_id_std.shape)
row, col = sorted_att_std
top_n = 100
att_std_idx = row[(top_n - 1)], col[(top_n - 1)]
att_top_std = att_trials_id[att_std_idx]
att_top_std = F.softmax(torch.tensor(att_top_std), dim=-1).numpy()
att_top_std_im = att_top_std.reshape(1, 20, 4, 7)
att_top_std_im = att_top_std_im[-1, -1, :, :]
def heatmap2d(arr: np.ndarray, alpha=0.5):
    ncolors = 256
    color_array = plt.get_cmap('viridis')(range(ncolors))

    # change alpha values
    n = 20
    color_array[:,-1] = [0.0] * n +  np.linspace(0.0,1.0,(ncolors - n)).tolist()

    # create a colormap object
    map_object = LinearSegmentedColormap.from_list(name='rainbow_alpha',colors=color_array)

    # register this new colormap with matplotlib
    # plt.clim(0, 1)
    plt.register_cmap(cmap=map_object)
    arr = gaussian_filter(arr, sigma=3.5)
    h = plt.imshow(arr, cmap='rainbow_alpha', alpha=alpha)
    # plt.colorbar()
    # plt.clim(0, 1)
    # plt.colorbar(mappable=h)
from utils import set_plot_white, set_plot_black
set_plot_black()
from scipy.ndimage.filters import gaussian_filter
from matplotlib.colors import LinearSegmentedColormap
from scipy.special import softmax

ix_step = [171, 197, 65, 217]
ncol = 2
nrow = len(ix_step) // ncol + 1
plt.figure(figsize=(20, 20))

sorted_att_std = np.unravel_index(np.argsort(-att_trials_id_std.ravel()), att_trials_id_std.shape)
row, col = sorted_att_std
for n, ix in enumerate(ix_step):
    top_n = ix
    att_std_idx = row[(top_n - 1)], col[(top_n - 1)]
    att_top_std = att_trials_id[att_std_idx]
    att_top_std = F.softmax(torch.tensor(att_top_std), dim=-1).numpy()
    att_top_std_im = att_top_std.reshape(1, 20, 4, 7)
    att_top_std_im = att_top_std_im[-1, -1, :, :]

    t = df.groupby(['Interval', 'Trial']).size().reset_index().iloc[att_std_idx[0]]
    if t['Trial'] <= 20: n_stim = 0
    elif t['Trial'] <= 40: n_stim = 1
    elif t['Trial'] <= 60: n_stim = 2  

    im_interval = vid_stack[n_stim, math.ceil(t['Interval'] * 20)]

    att_grid =  softmax((att_top_std_im / att_top_std_im.max())**1000)
    att_grid = np.repeat(att_grid, (128 // 8), axis=0)
    att_grid = np.repeat(att_grid, (128 // 8), axis=1)

    plt.subplot(nrow, ncol, n+1)
    # im_interval = gaussian_filter(im_interval, sigma=1)
    plt.imshow(im_interval, cmap='gray')
    heatmap2d(att_grid, alpha=0.4)
    plt.title(f'Stimulus Attention Heatmap {i}', fontsize=20)<jupyter_output>/tmp/ipykernel_8142/3847477116.py:14: UserWarning: Trying to register the cmap 'rainbow_alpha' which already exists.
  plt.register_cmap(cmap=map_object)
StarcoderdataJupyterStructuredCleanDedup