I'm new to PyTorch and I'm writing a unit test for an activation function I'm making.

I plan to test against a reference implementation for this function. I want to approach this in a test-driven way, so I learned to write a test using a known-good function: the ReLU implementation "MyReLU" from this beginner tutorial.

The tests passed, but is there any way I can improve the code below? I worry I might not be taking full advantage of PyTorch's libraries and capabilities.

import unittest
import numpy as np
import torch
from torch.autograd import gradcheck
from my_activation_functions import MyReLU

class ReluTest(unittest.TestCase):
    def setUp(self):
        self.relu = MyReLU.apply

    def test_relu_values_x_leqz(self):
        tin_leqz = torch.tensor(np.linspace(-10,0,300))
        tout_leqz = list(self.relu(tin_leqz))
        for x in tout_leqz:

    def test_relu_values_x0(self):
        tin_eqz = torch.tensor([0,0,0,0,0])
        tout_eqz = list(self.relu(tin_eqz))
        for x in tout_eqz:

    def test_relu_values_x_geqz(self):
        tin_geqz = torch.tensor(np.linspace(0.001,10,300))
        tout_geqz = list(self.relu(tin_geqz))
        test_geqz = list(tin_geqz)
        for ii in range(len(tout_geqz)):
            self.assertEqual(tout_geqz[ii], test_geqz[ii])

    def test_drelu_values(self):
        tin = (torch.randn(20,20,dtype=torch.double,requires_grad=True))
        self.assertTrue(gradcheck(self.relu, tin, eps=1e-6, atol=1e-4))

if __name__ == '__main__':

1 Answer 1


Since you are asking about PyTorch's capabilities you are not taking advantage of, you might want to use:

  • torch.linspace(-10,0,300) instead of torch.tensor(np.linspace(-10,0,300))

  • torch.zeros(5, dtype=torch.long) instead of torch.tensor([0,0,0,0,0])

  • tensor operations instead of iterating over each element of the tensor in a loop. This might not matter much in unit-tests but is important if you want to get GPU acceleration:

    self.assertTrue(torch.equal(tout_leqz, torch.zeros_like(tin_leqz)))


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.