|
- import unittest
- import os
- from itertools import chain
-
- import torch
- import paddle
- from paddle.io import Dataset, DataLoader
- import numpy as np
-
- from fastNLP.modules.mix_modules.mix_module import MixModule
- from fastNLP.modules.mix_modules.utils import paddle2torch, torch2paddle
- from fastNLP.core import synchronize_safe_rm
-
-
- ############################################################################
- #
- # 测试类的基本功能
- #
- ############################################################################
-
- class TestMixModule(MixModule):
- def __init__(self):
- super(TestMixModule, self).__init__()
-
- self.torch_fc1 = torch.nn.Linear(10, 10)
- self.torch_softmax = torch.nn.Softmax(0)
- self.torch_conv2d1 = torch.nn.Conv2d(10, 10, 3)
- self.torch_tensor = torch.ones(3, 3)
- self.torch_param = torch.nn.Parameter(torch.ones(4, 4))
-
- self.paddle_fc1 = paddle.nn.Linear(10, 10)
- self.paddle_softmax = paddle.nn.Softmax(0)
- self.paddle_conv2d1 = paddle.nn.Conv2D(10, 10, 3)
- self.paddle_tensor = paddle.ones((4, 4))
-
- class TestTorchModule(torch.nn.Module):
- def __init__(self):
- super(TestTorchModule, self).__init__()
-
- self.torch_fc1 = torch.nn.Linear(10, 10)
- self.torch_softmax = torch.nn.Softmax(0)
- self.torch_conv2d1 = torch.nn.Conv2d(10, 10, 3)
- self.torch_tensor = torch.ones(3, 3)
- self.torch_param = torch.nn.Parameter(torch.ones(4, 4))
-
- class TestPaddleModule(paddle.nn.Layer):
- def __init__(self):
- super(TestPaddleModule, self).__init__()
-
- self.paddle_fc1 = paddle.nn.Linear(10, 10)
- self.paddle_softmax = paddle.nn.Softmax(0)
- self.paddle_conv2d1 = paddle.nn.Conv2D(10, 10, 3)
- self.paddle_tensor = paddle.ones((4, 4))
-
-
- class TorchPaddleMixModuleTestCase(unittest.TestCase):
-
- def setUp(self):
-
- self.model = TestMixModule()
- self.torch_model = TestTorchModule()
- self.paddle_model = TestPaddleModule()
-
- def test_to(self):
- """
- 测试混合模型的to函数
- """
-
- self.model.to("cuda")
- self.torch_model.to("cuda")
- self.paddle_model.to("gpu")
- self.if_device_correct("cuda")
-
- self.model.to("cuda:2")
- self.torch_model.to("cuda:2")
- self.paddle_model.to("gpu:2")
- self.if_device_correct("cuda:2")
-
- self.model.to("gpu:1")
- self.torch_model.to("cuda:1")
- self.paddle_model.to("gpu:1")
- self.if_device_correct("cuda:1")
-
- self.model.to("cpu")
- self.torch_model.to("cpu")
- self.paddle_model.to("cpu")
- self.if_device_correct("cpu")
-
- def test_train_eval(self):
- """
- 测试train和eval函数
- """
-
- self.model.eval()
- self.if_training_correct(False)
-
- self.model.train()
- self.if_training_correct(True)
-
- def test_parameters(self):
- """
- 测试parameters()函数,由于初始化是随机的,目前仅比较得到结果的长度
- """
- mix_params = []
- params = []
-
- for value in self.model.named_parameters():
- mix_params.append(value)
-
- for value in chain(self.torch_model.named_parameters(), self.paddle_model.named_parameters()):
- params.append(value)
-
- self.assertEqual(len(params), len(mix_params))
-
- def test_named_parameters(self):
- """
- 测试named_parameters函数
- """
-
- mix_param_names = []
- param_names = []
-
- for name, value in self.model.named_parameters():
- mix_param_names.append(name)
-
- for name, value in chain(self.torch_model.named_parameters(), self.paddle_model.named_parameters()):
- param_names.append(name)
-
- self.assertListEqual(sorted(param_names), sorted(mix_param_names))
-
- def test_torch_named_parameters(self):
- """
- 测试对torch参数的提取
- """
-
- mix_param_names = []
- param_names = []
-
- for name, value in self.model.named_parameters(backend="torch"):
- mix_param_names.append(name)
-
- for name, value in self.torch_model.named_parameters():
- param_names.append(name)
-
- self.assertListEqual(sorted(param_names), sorted(mix_param_names))
-
- def test_paddle_named_parameters(self):
- """
- 测试对paddle参数的提取
- """
-
- mix_param_names = []
- param_names = []
-
- for name, value in self.model.named_parameters(backend="paddle"):
- mix_param_names.append(name)
-
- for name, value in self.paddle_model.named_parameters():
- param_names.append(name)
-
- self.assertListEqual(sorted(param_names), sorted(mix_param_names))
-
- def test_torch_state_dict(self):
- """
- 测试提取torch的state dict
- """
- torch_dict = self.torch_model.state_dict()
- mix_dict = self.model.state_dict(backend="torch")
-
- self.assertListEqual(sorted(torch_dict.keys()), sorted(mix_dict.keys()))
-
- def test_paddle_state_dict(self):
- """
- 测试提取paddle的state dict
- """
- paddle_dict = self.paddle_model.state_dict()
- mix_dict = self.model.state_dict(backend="paddle")
-
- # TODO 测试程序会显示passed后显示paddle的异常退出信息
- self.assertListEqual(sorted(paddle_dict.keys()), sorted(mix_dict.keys()))
-
- def test_state_dict(self):
- """
- 测试提取所有的state dict
- """
- all_dict = self.torch_model.state_dict()
- all_dict.update(self.paddle_model.state_dict())
- mix_dict = self.model.state_dict()
-
- # TODO 测试程序会显示passed后显示paddle的异常退出信息
- self.assertListEqual(sorted(all_dict.keys()), sorted(mix_dict.keys()))
-
- def test_load_state_dict(self):
- """
- 测试load_state_dict函数
- """
- state_dict = self.model.state_dict()
-
- new_model = TestMixModule()
- new_model.load_state_dict(state_dict)
- new_state_dict = new_model.state_dict()
-
- for name, value in state_dict.items():
- state_dict[name] = value.tolist()
- for name, value in new_state_dict.items():
- new_state_dict[name] = value.tolist()
-
- self.assertDictEqual(state_dict, new_state_dict)
-
- def test_save_and_load_state_dict(self):
- """
- 测试save_state_dict_to_file和load_state_dict_from_file函数
- """
- path = "model"
- try:
- self.model.save_state_dict_to_file(path)
- new_model = TestMixModule()
- new_model.load_state_dict_from_file(path)
-
- state_dict = self.model.state_dict()
- new_state_dict = new_model.state_dict()
-
- for name, value in state_dict.items():
- state_dict[name] = value.tolist()
- for name, value in new_state_dict.items():
- new_state_dict[name] = value.tolist()
-
- self.assertDictEqual(state_dict, new_state_dict)
- finally:
- synchronize_safe_rm(path)
-
- def if_device_correct(self, device):
-
-
- self.assertEqual(self.model.torch_fc1.weight.device, self.torch_model.torch_fc1.weight.device)
- self.assertEqual(self.model.torch_conv2d1.weight.device, self.torch_model.torch_fc1.bias.device)
- self.assertEqual(self.model.torch_conv2d1.bias.device, self.torch_model.torch_conv2d1.bias.device)
- self.assertEqual(self.model.torch_tensor.device, self.torch_model.torch_tensor.device)
- self.assertEqual(self.model.torch_param.device, self.torch_model.torch_param.device)
-
- if device == "cpu":
- self.assertTrue(self.model.paddle_fc1.weight.place.is_cpu_place())
- self.assertTrue(self.model.paddle_fc1.bias.place.is_cpu_place())
- self.assertTrue(self.model.paddle_conv2d1.weight.place.is_cpu_place())
- self.assertTrue(self.model.paddle_conv2d1.bias.place.is_cpu_place())
- self.assertTrue(self.model.paddle_tensor.place.is_cpu_place())
- elif device.startswith("cuda"):
- self.assertTrue(self.model.paddle_fc1.weight.place.is_gpu_place())
- self.assertTrue(self.model.paddle_fc1.bias.place.is_gpu_place())
- self.assertTrue(self.model.paddle_conv2d1.weight.place.is_gpu_place())
- self.assertTrue(self.model.paddle_conv2d1.bias.place.is_gpu_place())
- self.assertTrue(self.model.paddle_tensor.place.is_gpu_place())
-
- self.assertEqual(self.model.paddle_fc1.weight.place.gpu_device_id(), self.paddle_model.paddle_fc1.weight.place.gpu_device_id())
- self.assertEqual(self.model.paddle_fc1.bias.place.gpu_device_id(), self.paddle_model.paddle_fc1.bias.place.gpu_device_id())
- self.assertEqual(self.model.paddle_conv2d1.weight.place.gpu_device_id(), self.paddle_model.paddle_conv2d1.weight.place.gpu_device_id())
- self.assertEqual(self.model.paddle_conv2d1.bias.place.gpu_device_id(), self.paddle_model.paddle_conv2d1.bias.place.gpu_device_id())
- self.assertEqual(self.model.paddle_tensor.place.gpu_device_id(), self.paddle_model.paddle_tensor.place.gpu_device_id())
- else:
- raise NotImplementedError
-
- def if_training_correct(self, training):
-
- self.assertEqual(self.model.torch_fc1.training, training)
- self.assertEqual(self.model.torch_softmax.training, training)
- self.assertEqual(self.model.torch_conv2d1.training, training)
-
- self.assertEqual(self.model.paddle_fc1.training, training)
- self.assertEqual(self.model.paddle_softmax.training, training)
- self.assertEqual(self.model.paddle_conv2d1.training, training)
-
-
- ############################################################################
- #
- # 测试在MNIST数据集上的表现
- #
- ############################################################################
-
- class MNISTDataset(Dataset):
- def __init__(self, dataset):
-
- self.dataset = [
- (
- np.array(img).astype('float32').reshape(-1),
- label
- ) for img, label in dataset
- ]
-
- def __getitem__(self, idx):
- return self.dataset[idx]
-
- def __len__(self):
- return len(self.dataset)
-
- class MixMNISTModel(MixModule):
- def __init__(self):
- super(MixMNISTModel, self).__init__()
-
- self.fc1 = paddle.nn.Linear(784, 64)
- self.fc2 = paddle.nn.Linear(64, 32)
- self.fc3 = torch.nn.Linear(32, 10)
- self.fc4 = torch.nn.Linear(10, 10)
-
- def forward(self, x):
-
- paddle_out = self.fc1(x)
- paddle_out = self.fc2(paddle_out)
- torch_in = paddle2torch(paddle_out)
- torch_out = self.fc3(torch_in)
- torch_out = self.fc4(torch_out)
-
- return torch_out
-
- class TestMNIST(unittest.TestCase):
-
- @classmethod
- def setUpClass(self):
-
- self.train_dataset = paddle.vision.datasets.MNIST(mode='train')
- self.test_dataset = paddle.vision.datasets.MNIST(mode='test')
- self.train_dataset = MNISTDataset(self.train_dataset)
-
- self.lr = 0.0003
- self.epochs = 20
-
- self.dataloader = DataLoader(self.train_dataset, batch_size=100, shuffle=True)
-
- def setUp(self):
-
- self.model = MixMNISTModel().to("cuda")
- self.torch_loss_func = torch.nn.CrossEntropyLoss()
-
- self.torch_opt = torch.optim.Adam(self.model.parameters(backend="torch"), self.lr)
- self.paddle_opt = paddle.optimizer.Adam(parameters=self.model.parameters(backend="paddle"), learning_rate=self.lr)
-
- def test_case1(self):
-
- # 开始训练
- for epoch in range(self.epochs):
- epoch_loss, batch = 0, 0
- for batch, (img, label) in enumerate(self.dataloader):
-
- img = paddle.to_tensor(img).cuda()
- torch_out = self.model(img)
- label = torch.from_numpy(label.numpy()).reshape(-1)
- loss = self.torch_loss_func(torch_out.cpu(), label)
- epoch_loss += loss.item()
-
- loss.backward()
- self.torch_opt.step()
- self.paddle_opt.step()
- self.torch_opt.zero_grad()
- self.paddle_opt.clear_grad()
-
- else:
- self.assertLess(epoch_loss / (batch + 1), 0.3)
-
- # 开始测试
- correct = 0
- for img, label in self.test_dataset:
-
- img = paddle.to_tensor(np.array(img).astype('float32').reshape(1, -1))
- torch_out = self.model(img)
- res = torch_out.softmax(-1).argmax().item()
- label = label.item()
- if res == label:
- correct += 1
-
- acc = correct / len(self.test_dataset)
- self.assertGreater(acc, 0.85)
-
- ############################################################################
- #
- # 测试在ERNIE中文数据集上的表现
- #
- ############################################################################
|