Principle Directions from Weights.py

random_direction1 = []
random_direction2 = []

for w in copy_of_the_weights:
    
    if w.dim() == 1: 
        random_direction1.append(torch.zeros_like(w))
        random_direction2.append(torch.zeros_like(w))
        
    else:
        random_vector = w.clone().cpu().numpy()
        
        random_vector1 = random_vector - random_vector.mean((2,3),keepdims=True)
        random_vector2 = random_vector - random_vector.mean((0,1),keepdims=True)
        random_vector2 = np.transpose(random_vector2,(2,3,0,1))
        
        s1,u1,v1 = tf.linalg.svd(random_vector1,False)
        s2,u2,v2 = tf.linalg.svd(random_vector2,False)
        
        random_vector1 = u1 @ tf.linalg.diag(s1)[:,:,:,:1] @ tf.transpose(v1,(0,1,3,2))[:,:,:1,:]
        random_vector2 = u2 @ tf.linalg.diag(s2)[:,:,:,:1] @ tf.transpose(v2,(0,1,3,2))[:,:,:1,:]
        random_vector2 = tf.transpose(random_vector2,(2,3,0,1))
        
        random_vector1 = torch.from_numpy(random_vector1.eval()).cuda()
        random_vector2 = torch.from_numpy(random_vector2.eval()).cuda()

        w_norm  = w.view((w.shape[0],-1))  .norm(dim=(1),keepdim=True)[:,:,None,None]
        d_norm1 = random_vector1.view((random_vector1.shape[0],-1)).norm(dim=(1),keepdim=True)[:,:,None,None]
        d_norm2 = random_vector2.view((random_vector2.shape[0],-1)).norm(dim=(1),keepdim=True)[:,:,None,None]

        random_vector1 = random_vector1 * (w_norm/(d_norm1.cuda()+1e-10))
        random_vector2 = random_vector2 * (w_norm/(d_norm2.cuda()+1e-10))  
        
        print(random_vector1.shape)
        print(random_vector2.shape)

        random_direction1.append(random_vector1)
        random_direction2.append(random_vector2)