|
|
import unittest |
|
|
|
|
|
from apex import amp |
|
|
import random |
|
|
import torch |
|
|
from torch import nn |
|
|
|
|
|
from utils import common_init, HALF |
|
|
|
|
|
class TestRnnCells(unittest.TestCase): |
|
|
def setUp(self): |
|
|
self.handle = amp.init(enabled=True) |
|
|
common_init(self) |
|
|
|
|
|
def tearDown(self): |
|
|
self.handle._deactivate() |
|
|
|
|
|
def run_cell_test(self, cell, state_tuple=False): |
|
|
shape = (self.b, self.h) |
|
|
for typ in [torch.float, torch.half]: |
|
|
xs = [torch.randn(shape, dtype=typ).requires_grad_() |
|
|
for _ in range(self.t)] |
|
|
hidden_fn = lambda: torch.zeros(shape, dtype=typ) |
|
|
if state_tuple: |
|
|
hidden = (hidden_fn(), hidden_fn()) |
|
|
else: |
|
|
hidden = hidden_fn() |
|
|
outputs = [] |
|
|
for i in range(self.t): |
|
|
hidden = cell(xs[i], hidden) |
|
|
if state_tuple: |
|
|
output = hidden[0] |
|
|
else: |
|
|
output = hidden |
|
|
outputs.append(output) |
|
|
for y in outputs: |
|
|
self.assertEqual(y.type(), HALF) |
|
|
outputs[-1].float().sum().backward() |
|
|
for i, x in enumerate(xs): |
|
|
self.assertEqual(x.grad.dtype, x.dtype) |
|
|
|
|
|
def test_rnn_cell_is_half(self): |
|
|
cell = nn.RNNCell(self.h, self.h) |
|
|
self.run_cell_test(cell) |
|
|
|
|
|
def test_gru_cell_is_half(self): |
|
|
cell = nn.GRUCell(self.h, self.h) |
|
|
self.run_cell_test(cell) |
|
|
|
|
|
def test_lstm_cell_is_half(self): |
|
|
cell = nn.LSTMCell(self.h, self.h) |
|
|
self.run_cell_test(cell, state_tuple=True) |
|
|
|
|
|
class TestRnns(unittest.TestCase): |
|
|
def setUp(self): |
|
|
self.handle = amp.init(enabled=True) |
|
|
common_init(self) |
|
|
|
|
|
def tearDown(self): |
|
|
self.handle._deactivate() |
|
|
|
|
|
def run_rnn_test(self, rnn, layers, bidir, state_tuple=False): |
|
|
for typ in [torch.float, torch.half]: |
|
|
x = torch.randn((self.t, self.b, self.h), dtype=typ).requires_grad_() |
|
|
hidden_fn = lambda: torch.zeros((layers + (layers * bidir), |
|
|
self.b, self.h), dtype=typ) |
|
|
if state_tuple: |
|
|
hidden = (hidden_fn(), hidden_fn()) |
|
|
else: |
|
|
hidden = hidden_fn() |
|
|
output, _ = rnn(x, hidden) |
|
|
self.assertEqual(output.type(), HALF) |
|
|
output[-1, :, :].float().sum().backward() |
|
|
self.assertEqual(x.grad.dtype, x.dtype) |
|
|
|
|
|
def test_rnn_is_half(self): |
|
|
configs = [(1, False), (2, False), (2, True)] |
|
|
for layers, bidir in configs: |
|
|
rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=layers, |
|
|
nonlinearity='relu', bidirectional=bidir) |
|
|
self.run_rnn_test(rnn, layers, bidir) |
|
|
|
|
|
def test_gru_is_half(self): |
|
|
configs = [(1, False), (2, False), (2, True)] |
|
|
for layers, bidir in configs: |
|
|
rnn = nn.GRU(input_size=self.h, hidden_size=self.h, num_layers=layers, |
|
|
bidirectional=bidir) |
|
|
self.run_rnn_test(rnn, layers, bidir) |
|
|
|
|
|
def test_lstm_is_half(self): |
|
|
configs = [(1, False), (2, False), (2, True)] |
|
|
for layers, bidir in configs: |
|
|
rnn = nn.LSTM(input_size=self.h, hidden_size=self.h, num_layers=layers, |
|
|
bidirectional=bidir) |
|
|
self.run_rnn_test(rnn, layers, bidir, state_tuple=True) |
|
|
|
|
|
def test_rnn_packed_sequence(self): |
|
|
num_layers = 2 |
|
|
rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=num_layers) |
|
|
for typ in [torch.float, torch.half]: |
|
|
x = torch.randn((self.t, self.b, self.h), dtype=typ).requires_grad_() |
|
|
lens = sorted([random.randint(self.t // 2, self.t) for _ in range(self.b)], |
|
|
reverse=True) |
|
|
|
|
|
torch.set_default_tensor_type(torch.FloatTensor) |
|
|
lens = torch.tensor(lens, dtype=torch.int64, device=torch.device('cpu')) |
|
|
packed_seq = nn.utils.rnn.pack_padded_sequence(x, lens) |
|
|
torch.set_default_tensor_type(torch.cuda.FloatTensor) |
|
|
hidden = torch.zeros((num_layers, self.b, self.h), dtype=typ) |
|
|
output, _ = rnn(packed_seq, hidden) |
|
|
self.assertEqual(output.data.type(), HALF) |
|
|
output.data.float().sum().backward() |
|
|
self.assertEqual(x.grad.dtype, x.dtype) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
unittest.main() |
|
|
|
