Reset a model parameter value regularly during training

[anupsingh15]

I have created a model in the following way. The model mimics the straight-through gradient estimation as proposed in the VQ-VAE, with the difference that I reset the codewords parameter after every epoch for some training epochs.

class KMeansQuantizer(nn.Module):
    def __init__(self, num_codewords: int, dim: int, seglen: int, filespath: str, codewords_trainable: Optional[bool]=False):
        super().__init__()
        self.num_codewords = num_codewords
        self.dim = dim
        self.codewords_trainable = codewords_trainable
        if not self.codewords_trainable:
            self.db = (self.get_segs(filespath, seglen))
        self.codewords = nn.Parameter(torch.ones(num_codewords, dim)) #nn.Parameter(0.01*torch.randn(num_codewords, dim))

    @staticmethod
    def get_segs(filespath: str, seglen: int)-> torch.Tensor:
        # some steps
        return db

    def run_kmeans(self, model: List[nn.Module], niter:Optional[int]=100, verbose:Optional[bool]=True, seed:Optional[int]=100, device=Optional[str])-> None:
        model[0].eval()
        model[1].eval()
        Z_db = []
        with torch.no_grad():  
            print("preparing database for clustering")
            for batch in tqdm(self.db):
                ze = model[1](model[0](batch.to(device)))
                Z_db.append(ze.detach().cpu().numpy())
            Z_db = np.concatenate(Z_db).reshape(-1, self.dim)
            self.kmeans = faiss.Kmeans(self.dim, self.num_codewords, niter=niter, verbose=verbose, spherical=True)
            self.kmeans.seed = seed
            self.kmeans.train(Z_db)  
        self.codewords.data =  torch.from_numpy(self.kmeans.centroids).to(device).detach().clone() ######
        model[0].train()
        model[1].train()
        Z_db = None
        return 

    def _pass_grad(self, x: torch.Tensor, y:torch.Tensor)-> torch.Tensor:
        """
        performs gradients bypassing
        for y = f(x), where f is non differentiable, we set dL/dy = dL/dx
        """
        return y.detach() + (x-x.detach())
    
    @staticmethod
    def _get_cosine_disimilarity(x,y):
        cosine_dist = 1 - torch.einsum('bd, bd -> b', [x.view(-1, x.shape[-1]), y.view(-1, y.shape[-1])])
        return cosine_dist.mean()
        
    def forward(self, ze):
        results = {}
        if self.codewords_trainable:
            codewords = self.codewords/torch.linalg.norm(self.codewords, dim=-1).view(-1,1)
        else:
            codewords = self.codewords
        results["codewords"] = codewords
        
        nearest_clust_idx = torch.argmax(torch.einsum("btd,kd -> btk", [ze, codewords]), dim=-1)
        zq = codewords[nearest_clust_idx]
        x = self._pass_grad(ze, zq)
        results["idx"] = nearest_clust_idx

        if self.codewords_trainable:
            results["latent_loss"] = self._get_cosine_disimilarity(zq, ze.detach())
            results["commitment_loss"] = self._get_cosine_disimilarity(ze, zq.detach())
        

I want to reset the parameter value before a training epoch starts. It will get updated with every training step, but I'd like to reset its values for the next training epoch. I want to reiterate this process for some training epochs only.

The way I do it right now is to reset its data value (in the run_kmeans method) as mentioned hereunder, but I get its gradient to be None:

self.codewords.data = torch.from_numpy(self.kmeans.centroids).to(device).detach().clone()

Could someone please let me know what is the correct implementation?