hydra_init(cfg_name) # 파라미터 설정
hydra_main() # train 수행

distributed_utils.call_main(cfg, pre_main)

distributed_main(i, main, cfg: FairseqConfig, kwargs)

features = self.feature_extractor(source)


def forward(self, x):
    ''''
    input: x.shape
    torch.Size([8, 101168])

    # output : x.shape 
    torch.Size([8, 512, 315])
    '''
    # BxT -> BxCxT
    x = x.unsqueeze(1)

    for conv in self.conv_layers:
        x = conv(x)

    return x

if self.feature_grad_mult > 0:
    features = self.feature_extractor(source)
    if self.feature_grad_mult != 1.0:
        '''
        gradiant scaling
        '''
        features = GradMultiply.apply(features, self.feature_grad_mult)

# torch.Size([8, 512, 315]) => torch.Size([8, 315, 512])
features = features.transpose(1, 2)
features = self.layer_norm(features)
unmasked_features = features.clone()

if self.post_extract_proj is not None:
    '''
    features.shape
    torch.Size([8, 315, 512]) => torch.Size([8, 315, 768])
    '''
    features = self.post_extract_proj(features)

features = self.dropout_input(features)
unmasked_features = self.dropout_features(unmasked_features)

def apply_mask(self, features, padding_mask):
    if self.mask_prob > 0:
        mask_indices = compute_mask_indices(
            shape = (8, 315),
            padding_mask = None,
            mask_prob = 0.65,
            mask_length = 10,
            mask_type = 'static', # fixed size 
            mask_other = 0.0,
            min_masks=2, # minimum number of masked spans
            no_overlap=False, #if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
            min_space=1, # only used if no_overlap is True, this is how many elements to keep unmasked between spans
        )
        mask_indices = torch.from_numpy(mask_indices).to(x.device)
        x[mask_indices] = self.mask_emb

'''
compute_mask_indices function 핵심 코드 부분
'''
all_num_mask = int(
    # add a random number for probabilistic rounding
    mask_prob * all_sz / float(mask_length)
    + np.random.rand()
)
all_num_mask = max(min_masks, all_num_mask)

if mask_type == "static":
    lengths = np.full(num_mask, mask_length)
'''
lengths
array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
    10, 10, 10])
'''

mask_idc = np.asarray(
    [
        mask_idc[j] + offset
        for j in range(len(mask_idc))
        for offset in range(lengths[j])
    ]
)
mask_idc
array([ 80,  81,  82,  83,  84,  85,  86,  87,  88,  89, 298, 299, 300,
       301, 302, 303, 304, 305, 306, 307,  72,  73,  74,  75,  76,  77,
        78,  79,  80,  81, 253, 254, 255, 256, 257, 258, 259, 260, 261,
       262, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227,  89,  90,
        91,  92,  93,  94,  95,  96,  97,  98,  37,  38,  39,  40,  41,
        42,  43,  44,  45,  46, 228, 229, 230, 231, 232, 233, 234, 235,
       236, 237, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 287,
       288, 289, 290, 291, 292, 293, 294, 295, 296, 224, 225, 226, 227,
       228, 229, 230, 231, 232, 233, 240, 241, 242, 243, 244, 245, 246,
       247, 248, 249,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,
        59,  60,  61,  62,  63,  64,  65,  66,  67,  68, 185, 186, 187,
       188, 189, 190, 191, 192, 193, 194, 268, 269, 270, 271, 272, 273,
       274, 275, 276, 277, 164, 165, 166, 167, 168, 169, 170, 171, 172,
       173, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 304, 305,
       306, 307, 308, 309, 310, 311, 312, 313, 112, 113, 114, 115, 116,
       117, 118, 119, 120, 121])

np.unique(mask_idc[mask_idc < sz])
array([ 37,  38,  39,  40,  41,  42,  43,  44,  45,  46,  58,  59,  60,
        61,  62,  63,  64,  65,  66,  67,  68,  72,  73,  74,  75,  76,
        77,  78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,
        90,  91,  92,  93,  94,  95,  96,  97,  98, 102, 103, 104, 105,
       106, 107, 108, 109, 110, 111, 112, 113, 114, 123, 124, 125, 126,
       127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
       140, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 187, 188,
       189, 190, 191, 192, 193, 194, 195, 196, 221, 222, 223, 224, 225,
       226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238,
       239, 240, 241, 242, 243, 257, 258, 259, 260, 261, 262, 263, 264,
       265, 266, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 288,
       289, 290, 291, 292, 293, 294, 295, 296, 297, 303, 304, 305, 306,
       307, 308, 309, 310, 311, 312])

min_len = min([len(m) for m in mask_idcs])
for i, mask_idc in enumerate(mask_idcs):
    if len(mask_idc) > min_len:
        mask_idc = np.random.choice(mask_idc, min_len, replace=False)
    mask[i, mask_idc] = True

x, mask_indices = self.apply_mask(features, padding_mask)

x.shape # float
torch.Size([8, 315, 768])

mask_indices.shape # bool
torch.Size([8, 315])

x = self.encoder(x, padding_mask=None)

# in self.encoder
'''
x.shape
torch.Size([8, 315, 768])
'''
x = self.extract_features(x, padding_mask)

def extract_features(self, x, padding_mask=None):

    if padding_mask is not None:
        x[padding_mask] = 0
    # x: torch.Size([8, 315, 768]) => torch.Size([8, 768, 315])
    # x_conv: torch.Size([8, 768, 315])
    x_conv = self.pos_conv(x.transpose(1, 2))
    x_conv = x_conv.transpose(1, 2)
    x += x_conv

    if not self.layer_norm_first:
        x = self.layer_norm(x)

    x = F.dropout(x, p=self.dropout, training=self.training)

    # B x T x C -> T x B x C
    x = x.transpose(0, 1)

    layer_results = []
    for i, layer in enumerate(self.layers):
        dropout_probability = np.random.random()
        if not self.training or (dropout_probability > self.layerdrop):
            x, z = layer(x, self_attn_padding_mask=padding_mask, need_weights=False)
            layer_results.append(x)

    # T x B x C -> B x T x C
    x = x.transpose(0, 1)

    return x

'''
self.quantizer
GumbelVectorQuantizer(
  (weight_proj): Linear(in_features=512, out_features=640, bias=True)
)

y.shape
torch.Size([8, 141, 512]) # target vector in the masked indice
'''
q = self.quantizer(y, produce_targets=False)

def forward(self, x=y, produce_targets=False):
    bsz, tsz, fsz = x.shape # torch.Size([8(bz), 141(seq), 512(in_dim)]) # orignal vec but be cloned
    x = x.reshape(-1, fsz) # torch.Size([1128(seq*bz), 512(in_dim)])
    x = self.weight_proj(x) # torch.Size([1128(seq*bz), 640(num_var *n_group)])
    x = x.view(bsz * tsz * self.groups, -1) # torch.Size([2256(seq*bz*n_group), 320(num_var )])
    _, k = x.max(-1) # 각 group axis에서 큰 값의 위치를 추출

hard_x = (
    x.new_zeros(*x.shape) # torch.Size([2256, 320])
    .scatter_(-1, k.view(-1, 1), 1.0) # 마지막 축에 대하여 k indice에 해당하는 부분에 1.0을 할당
    .view(bsz * tsz, self.groups, -1) # reshape
)
hard_x.shape 
torch.Size([1128, 2, 320]) # (seq*bz, n_group, num_var )

hard_probs = torch.mean(hard_x.float(), dim=0)

hard_probs.shape 
torch.Size([2, 320])

result["code_perplexity"] = torch.exp(
        -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
    ).sum()

avg_probs = torch.softmax(
    x.view(bsz * tsz, self.groups, -1).float(), dim=-1
).mean(dim=0) # torch.Size([2, 320])
result["prob_perplexity"] = torch.exp(
    -torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)
).sum()

x = x.view(bsz * tsz, -1) # torch.Size([1128(bz*seq), 640(n_group*num_var )])
vars = nn.Parameter(torch.FloatTensor(1, num_groups * num_vars, var_dim))

'''
x.unsqueeze(-1).shape
torch.Size([1128, 640, 1])

vars.shape # initialized
torch.Size([1, 640, 128])
'''
x = x.unsqueeze(-1) * vars # torch.Size([1128, 640, 128])
x = x.view(bsz * tsz, self.groups, self.num_vars, -1) # torch.Size([1128(bz*seq), 2(n_group), 320(num_var), 128(var_dim)])
x = x.sum(-2) # torch.Size([1128(bz*seq), 2(n_group), 128(var_dim)])
x = x.view(bsz, tsz, -1) # torch.Size([8(bz), 141(seq), 256(n_group*var_dim)])

y= q["x"]
y = self.project_q(y)

y.shape
torch.Size([8, 141, 256])

num = y.size(1) # 141

bsz, tsz, fsz = y.shape
y = y.view(-1, fsz)  # BTC => (BxT)C # torch.Size([1128, 256])

if self.n_negatives > 0:
    tszs = (
        buffered_arange(num) # torch.Size([141])
        .unsqueeze(-1) # torch.Size([141, 1])
        .expand(-1, self.n_negatives) # torch.Size([141, 100]), hstack
        .flatten() # torch.Size([14100])
    )

    neg_idxs = torch.randint(
        low=0, high=high - 1, size=(bsz, self.n_negatives * num) 
    ) #  torch.Size([8, 14100])
    '''
    tensor([[105,  59,   4,  ...,  69,  34, 117],
        [ 78, 132,  35,  ...,  17, 125,  58],
        [ 92, 101,  74,  ..., 100,  64, 106],
        ...,
        [ 45,  10, 128,  ...,  68,  95,  36],
        [ 31,  82,  96,  ...,  46,  42,  69],
        [115,  67, 131,  ...,   3, 108,  33]])
    '''
    neg_idxs[neg_idxs >= tszs] += 1 # torch.Size([8, 14100]) # 왼쪽에 더 (음성이시작되는 부분에) negative sample이 되도록 수정
    '''
    tensor([[106,  60,   5,  ...,  69,  34, 117],
            [ 79, 133,  36,  ...,  17, 125,  58],
            [ 93, 102,  75,  ..., 100,  64, 106],
            ...,
            [ 46,  11, 129,  ...,  68,  95,  36],
            [ 32,  83,  97,  ...,  46,  42,  69],
            [116,  68, 132,  ...,   3, 108,  33]])
    '''
    negs = y[neg_idxs.view(-1)] # torch.Size([112800, 256]) # negative samples seleted
    negs = negs.view( #     torch.Size([8, 141, 100, 256])
        bsz, num, self.n_negatives + self.cross_sample_negatives, fsz
    ).permute(
        2, 0, 1, 3
    )  # to torch.Size([100, 8, 141, 256])

mask_indices.shape
(8, 315) #(bz, audio_seq)

# embedding vector에서 mask된 부분을 추출
x = x[mask_indices].view(x.size(0), -1, x.size(-1))
x.shape # torch.Size([8, 141, 768])

self.final_proj
Linear(in_features=768, out_features=256, bias=True)
x = self.final_proj(x)
x.shape # torch.Size([8, 141, 256])

'''
x.shape: torch.Size([8, 141, 256]) # masked original vectors
y.shape: torch.Size([8, 141, 256]) # target vectors in the mask locations
negs.shape : torch.Size([100, 8, 141, 256]) #  100개 neg 샘플링 in 141 masked range 
'''
x = self.compute_preds(x, y, negs)

def compute_preds(self, x, y, negatives):

    neg_is_pos = (y == negatives).all(-1) # (y == negatives).shape : torch.Size([100, 8, 141, 256])
    y = y.unsqueeze(0) # torch.Size([1, 8, 141, 256])
    targets = torch.cat([y, negatives], dim=0) # torch.Size([101, 8, 141, 256])

    logits = torch.cosine_similarity(x.float(), targets.float(), dim=-1).type_as(x) # torch.Size([101, 8, 141])

    logits /= self.logit_temp

    if neg_is_pos.any():
        logits[1:][neg_is_pos] = float("-inf") 

    return logits # torch.Size([101, 8, 141])

logits = model.get_logits(net_output).float()

def get_logits(self, net_output):
    logits = net_output["x"] # torch.Size([101, 8, 141])
    logits = logits.transpose(0, 2) # torch.Size([141, 8, 101])
    logits = logits.reshape(-1, logits.size(-1)) # torch.Size([1128, 101])
    return logits

target = model.get_targets(sample, net_output)
def get_targets(self, sample, net_output, expand_steps=True):
    x = net_output["x"] # torch.Size([101, 8, 141])
    return x.new_zeros(x.size(1) * x.size(2), dtype=torch.long) # torch.Size([1128]), 0th true label is located

if self.infonce:
    loss = F.cross_entropy(
        logits,
        target,
        reduction="sum" if reduce else "none",
    )

def get_extra_losses(self, net_output):
    pen = []

    if "prob_perplexity" in net_output:
        pen.append(
            (net_output["num_vars"] - net_output["prob_perplexity"])
            / net_output["num_vars"]
        )

    if "features_pen" in net_output:
        pen.append(net_output["features_pen"])  # x^2

    return pen

for p, coef in zip(extra_losses, self.loss_weights):
    if coef != 0 and p is not None:
        p = coef * p.float() * sample_size
        loss += p
        losses.append(p)

# features: torch.Size([8, 315, 512])
features_pen = features.float().pow(2).mean()

logging_output = {
    "loss": loss.item() if reduce else loss, # cross_entropy, prob_perplexity, feature_pen
    "ntokens": sample_size,
    "nsentences": sample["id"].numel(),
    "sample_size": sample_size,
}