|
|
|
@ -14,6 +14,7 @@
|
|
|
|
|
|
|
|
|
|
from __future__ import print_function
|
|
|
|
|
import numpy as np
|
|
|
|
|
import paddle
|
|
|
|
|
import paddle.fluid as fluid
|
|
|
|
|
from paddle.fluid.layers.utils import flatten
|
|
|
|
|
from paddle.fluid.dygraph import declarative, ProgramTranslator
|
|
|
|
@ -151,5 +152,33 @@ class TestWithTrainAndEval(unittest.TestCase):
|
|
|
|
|
partial_layer._train_program)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class GPT2LMHeadModel(fluid.dygraph.Layer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
super(GPT2LMHeadModel, self).__init__()
|
|
|
|
|
self.embedding0 = paddle.nn.Embedding(20, 16)
|
|
|
|
|
self.embedding1 = paddle.nn.Embedding(20, 32)
|
|
|
|
|
self.lm_head_weight = paddle.to_tensor(
|
|
|
|
|
np.random.rand(2, 3).astype('float32'))
|
|
|
|
|
|
|
|
|
|
@declarative
|
|
|
|
|
def forward(self, x):
|
|
|
|
|
x = fluid.layers.reshape(x, shape=[-1, 6])
|
|
|
|
|
x1, x2, x3 = fluid.layers.split(input=x, dim=1, num_or_sections=3)
|
|
|
|
|
return x1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestPruneUnusedParamInProgram(unittest.TestCase):
|
|
|
|
|
def test_prune(self):
|
|
|
|
|
input_ids = np.array([[15, 11, 6, 3, 18, 13]]).astype("float32")
|
|
|
|
|
|
|
|
|
|
place = fluid.CPUPlace()
|
|
|
|
|
with fluid.dygraph.guard(place):
|
|
|
|
|
model = GPT2LMHeadModel()
|
|
|
|
|
model.eval()
|
|
|
|
|
input_ids = paddle.to_tensor(input_ids)
|
|
|
|
|
out = model(input_ids)
|
|
|
|
|
self.assertTrue(np.array_equal(out.numpy(), [[15, 11]]))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
unittest.main()
|
|
|
|
|