Example codes of IntRA-KD

[cardwing]

The example codes of IntRA-KD are given as follows. The following implementation is slightly different from the original paper and may cause numerical instability in some cases. You may need to carefully select the value of several hyperparameters so that IntRA-KD can work normally. If you have limited time and resources, you can just follow the same setting in our paper and use our reported value.

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import torch.nn as nn

#######AOI generation#######
#######You need to preprocess the original label to obtain the AOI map#########
#######Suppose M is the AOI map (h x w x n), F is the feature map (h_f x w_f x c), n is the number of classes, c is the number of channels############

smoothed_map = cv2.blur(M, (11, 11)) # kernel_size is a tunable hyperparameter
smoothed_map = cv2.resize(smoothed_map, (h_f, w_f))
mu_s = np.zeros((3, n, c))
mu_t = np.zeros((3, n, c))
cons = 1.0/(h_f*w_f) # we do not count the number of non-zero elements in AOI map as it requires additional computation and brings no more gains, one trick is that we more frequently sample areas that have road markings
small_num = 1e-10

#######moment pooling######
for i in range(h_f):
    for j in range(w_f):
        for cls_num in range(n):
            for channel in range(c):
                mu_s[0, cls_num, channel] = cons*np.sum(M[:, :, cls_num]*F[:, :, channel])
                mat_1 = M[:, :, cls_num]*F[:, :, channel]-mu_s[0, cls_num, channel]
                mu_s[1, cls_num, channel] = cons*np.sum(mat_1**2)
                mat_2 = (M[:, :, cls_num]*F[:, :, channel]-mu_s[0, cls_num, channel])/(mu_s[1, cls_num, channel]+small_num)
                mu_s[2, cls_num, channel] = cons*np.sum(mat_2**3)
########mu_t is calculated in a similar way#######

#######Inter-region affinity graph generation########
c_s_0 = cosine_similarity(mu_s[0, :, :], mu_s[0, :, :]) #c_t_0, c_s_1, ... are calculated similarly
loss = nn.MSELoss()
loss_intra_kd_0 = loss(c_s_0, c_t_0.detach())