forked from ccf-ai-infra/GPUCodeForces
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
# embeddingbag_torch.py
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
BATCH_SIZE = 1024
|
|
EMB_DIM = 256
|
|
VOCAB_SIZE = 50000
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, embedding_weight):
|
|
super().__init__()
|
|
self.criterion = nn.EmbeddingBag(
|
|
VOCAB_SIZE, EMB_DIM, mode='mean', sparse=False, _weight=embedding_weight
|
|
)
|
|
|
|
|
|
def forward(self, input: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
|
|
return self.criterion(input, offsets)
|
|
|
|
def get_inputs():
|
|
AVG_LEN = 10
|
|
TOTAL_INDICES = BATCH_SIZE * AVG_LEN
|
|
indices = torch.randint(0, VOCAB_SIZE, (TOTAL_INDICES,), dtype=torch.long)
|
|
offsets = [0]
|
|
current_offset = 0
|
|
for _ in range(BATCH_SIZE - 1):
|
|
seq_len = torch.randint(1, AVG_LEN * 2, (1,)).item()
|
|
current_offset += seq_len
|
|
offsets.append(current_offset)
|
|
|
|
offsets = torch.tensor(offsets, dtype=torch.long)
|
|
|
|
if offsets[-1].item() > indices.shape[0]:
|
|
indices = indices[:offsets[-1].item()]
|
|
|
|
return [indices.cuda(), offsets.cuda()]
|
|
|
|
def get_init_inputs():
|
|
weight = torch.randn(VOCAB_SIZE, EMB_DIM, dtype=torch.float32)
|
|
return [weight] |