Horrible Performance - Two Towers Model

PopHSKE: <p>I have been beating my head against the wall and come to the group to see if there is anything that looks like a culprit of very poor performance in a simple Two-Tower model. When i look at a hold out (future time period) set of customer -product interactions, the hit rate is far worse than just random products being offered. I am looking at evaluating at a top 5 product suggestion for the offline evaluation.</p>
<p>The model is very simple and is just a dot product of the cust and product embeddings. There are around 5,000 products and 5.6 million customers:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class TwoTowerMF(nn.Module):
    
    def __init__(self, n_items, n_users, emb_dim):
        super(TwoTowerMF, self).__init__()
        
        self.emb_dim = emb_dim
        self.item_emb = nn.Embedding(num_embeddings=n_items, embedding_dim=self.emb_dim, padding_idx=0)
        self.user_emb = nn.Embedding(num_embeddings=n_users, embedding_dim=self.emb_dim, padding_idx=0)
         
        self.dot = torch.matmul

        
    def forward(self, items, users):
        
        item_emb = self.item_emb(items)
        user_emb = self.user_emb(users)

        score = self.dot(user_emb, item_emb.t())
        
        return score
</code></pre>
<p>The dataset uses the interactions held in memory. The item ID, customer ID and CSP which is just the probability that the item would be in the batch is yielded :</p>
<pre><code class="lang-auto">class interactionDataset(Dataset):
    def __init__(self, item_user_inters):
        self.inters = torch.LongTensor(item_user_inters[:,0:2]) #  user and item
        self.csp = torch.Tensor(item_user_inters[:,2]) # csp
        
    def __len__(self):
        return len(self.inters)  # number interactions

    def __getitem__(self, idx):
        inter = self.inters[idx]
        csp = self.csp[idx]
        
        return inter[0], inter[1], csp # first array are item ID, second are user ID, third is csp
</code></pre>
<p>Im using DDP with a 4 GPU instance</p>
<pre><code>train_ds = interactionDataset(pdf_ints_train) 
</code></pre>
<pre><code class="lang-auto">if args.is_distributed:
        # distributed training sampler
        train_sampler = torch.utils.data.distributed.DistributedSampler(
                        train_ds, num_replicas=args.world_size, rank=rank
        )
    

    train_loader = DataLoader(train_ds, 
                              shuffle= not args.is_distributed, 
                              sampler= train_sampler if args.is_distributed else None,
                              batch_size=args.batch_size_train,
                              drop_last=False, 
                              num_workers = args.num_workers)
</code></pre>
<p>The training loop is straight forward (I have tried with and without the adjustment to the loss based on the csp).</p>
<pre><code class="lang-auto">def train(args, model, device, train_loader, optimizer, epoch, loss_fn, rank):
    
    model.train()
    train_loss = 0.0
    
    
    b_sz = len(next(iter(train_loader))[0])
    print(f"[GPU{args.rank}] Epoch {epoch} | Batchsize: {b_sz}")
        
    
    for batch_idx, (items, custs, csp) in enumerate(train_loader):
        custs = custs.to(device)
        items = items.to(device)
        csp = csp.to(device)


        optimizer.zero_grad()

        # Forward pass
        scores = model(items, custs)

        # Compute Loss
        loss = loss_fn(scores - torch.log(csp), torch.eye(custs.shape[0]).to(device))
        train_loss += loss.item()


        # Backward pass
        loss.backward()
        
        optimizer.step()
</code></pre>
<p>The loss is in-batch soft max which i think is setup just like Tensorflow recommenders.</p>
<pre><code class="lang-auto">model = TwoTowerMF(n_items=n_items+1, n_users=n_users+1, emb_dim=args.emb_dim)
    
    # must come before call to DDP
    model.to(device)
    
    if args.is_distributed:
        model = DDP(model,device_ids=[local_rank])
    
    # Pin each GPU to a single library process
    torch.cuda.set_device(local_rank)
    

    
    loss_fn = nn.CrossEntropyLoss(reduction='sum')
    optimizer = torch.optim.Adam(model.parameters(), lr = args.lr * math.sqrt(args.num_gpus))
</code></pre>
<p>I have tried various learning rates, batch sizes etc and the model is always horrible.</p>
<p>Looking the loss for the first epoch, the model decreases loss in each batch but rough hits a period where it stops much improvement.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/e/fe0aae08c0c838328ad664ab3acc6077ea3fb53a.png" data-download-href="https://discuss.pytorch.org/uploads/default/fe0aae08c0c838328ad664ab3acc6077ea3fb53a" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/e/fe0aae08c0c838328ad664ab3acc6077ea3fb53a.png" alt="image" data-base62-sha1="Afm9QPxXq9BW566r7V88CcKWwvE" width="267" height="500" data-dominant-color="E7F1FE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">427×799 11.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>