forked from ccf-ai-infra/GPUCodeForces
68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class SwishTorchModel(nn.Module):
|
|
"""PyTorch原生Swish激活函数实现"""
|
|
|
|
def __init__(self):
|
|
super(SwishTorchModel, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return x * torch.sigmoid(x)
|
|
|
|
def get_init_inputs():
|
|
"""获取模型初始化参数"""
|
|
return []
|
|
|
|
def get_inputs():
|
|
"""获取模型输入数据"""
|
|
torch.manual_seed(42)
|
|
return [torch.randn(128, 256)]
|
|
|
|
def test_swish():
|
|
"""测试Swish激活函数"""
|
|
print("=" * 60)
|
|
print("Swish激活函数测试")
|
|
print("=" * 60)
|
|
|
|
# 创建测试数据
|
|
torch.manual_seed(42)
|
|
input_tensor = torch.randn(128, 256)
|
|
|
|
# 创建模型
|
|
model = SwishTorchModel()
|
|
|
|
# 前向传播
|
|
output = model(input_tensor)
|
|
|
|
print(f"输入形状: {input_tensor.shape}")
|
|
print(f"输出形状: {output.shape}")
|
|
print(f"输入范围: [{input_tensor.min().item():.3f}, {input_tensor.max().item():.3f}]")
|
|
print(f"输出范围: [{output.min().item():.3f}, {output.max().item():.3f}]")
|
|
|
|
# 验证Swish函数
|
|
expected_output = input_tensor * torch.sigmoid(input_tensor)
|
|
diff = torch.abs(output - expected_output).max().item()
|
|
print(f"与预期输出最大差异: {diff:.6f}")
|
|
|
|
if diff < 1e-6:
|
|
print("✅ Swish激活函数实现正确")
|
|
else:
|
|
print("❌ Swish激活函数实现有误")
|
|
|
|
return output
|
|
|
|
def main():
|
|
"""主测试函数"""
|
|
try:
|
|
output = test_swish()
|
|
print("\n" + "=" * 60)
|
|
print("测试总结")
|
|
print("=" * 60)
|
|
print("✅ Swish激活函数测试完成")
|
|
|
|
except Exception as e:
|
|
print(f"测试过程中出现错误: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |