exploding_gradients.py

import torch
import torch.autograd
import torch.functional
import torch.nn
import torch.nn.functional as F  # noqa
import torch.optim
import torch.utils.data
import tqdm

import config
import experiment_buddy


def main(logger):
    torch.manual_seed(config.random_seed)

    loss = torch.nn.CrossEntropyLoss()
    p = torch.nn.Parameter(torch.zeros(2, 2))
    optimizer = torch.optim.SGD([p, ], lr=0.1)
    y = torch.tensor([0, 1])

    for epoch in tqdm.trange(config.num_epochs):
        l = loss(p, y)
        l.backward()
        print(l.item(), p.grad.flatten(), p.flatten().data)
        optimizer.step()
    print("done")


if __name__ == '__main__':
    experiment_buddy.register_defaults(vars(config))
    tb = experiment_buddy.deploy(host="")
    main(tb)