INNER CODE UNIT · Python

main

NVIDIA/DeepLearningExamples · PyTorch/LanguageModeling/BERT/lamb_amp_opt/main.py:37

def main() :
    loss = torch.nn.CrossEntropyLoss(ignore_index=-1)
    model = TestMod()
    model.cuda()
    model.half()
    model.train()
    model = torch.jit.script(model)
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'gamma', 'beta', 'LayerNorm']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
    grad_scaler = torch.cuda.amp.GradScaler(enabled=True)
    optimizer = FusedLAMBAMP(optimizer_grouped_parameters)

    x, y = make_classification(n_samples=N_SAMPLES, n_features=N_FEATURES, random_state=0)
    x = StandardScaler().fit_transform(x)
    inputs = torch.from_numpy(x).cuda().half()

View source record →