INNER CODE UNIT · Python

print_param_diff

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

def print_param_diff(optimizer):
    with torch.no_grad():
        for i, (group, master_group) in enumerate(zip(optimizer.param_groups, optimizer.param_groups_fp32)):
            for ii, (p, master_p) in enumerate(zip(group['params'], master_group['params'])):
                diff = (p - master_p.half()).float().abs().mean().item()
                print(f"  {i}th group, {ii}th param diff: {diff}")


class TestMod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(N_FEATURES, N_FEATURES // 2),
            torch.nn.ReLU(),
            torch.nn.Linear(N_FEATURES // 2, 2),
        )
    def forward(self, inputs) :
        return self.layers(inputs)

View source record →