docker-python/tests/test_transformers.py at master · AIComputerVision/docker-python · GitHub

25 lines (21 loc) · 929 Bytes

import unittest

import torch

from transformers import AdamW

class TestTransformers(unittest.TestCase):

def assertListAlmostEqual(self, list1, list2, tol):

self.assertEqual(len(list1), len(list2))

for a, b in zip(list1, list2):

self.assertAlmostEqual(a, b, delta=tol)

def test_adam_w(self):

w = torch.tensor([0.1, -0.2, -0.1], requires_grad=True)

target = torch.tensor([0.4, 0.2, -0.5])

criterion = torch.nn.MSELoss()

# No warmup, constant schedule, no gradient clipping

optimizer = AdamW(params=[w], lr=2e-1, weight_decay=0.0)

for _ in range(100):

loss = criterion(w, target)

loss.backward()

optimizer.step()

w.grad.detach_() # No zero_grad() function on simple tensors. we do it ourselves.

w.grad.zero_()

self.assertListAlmostEqual(w.tolist(), [0.4, 0.2, -0.5], tol=1e-2)

Expand file tree