forked from ccf-ai-infra/GPUCodeForces
81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class GELUDropoutTorchModel(nn.Module):
|
|
"""PyTorch原生GELU+Dropout实现"""
|
|
|
|
def __init__(self, p=0.1):
|
|
super(GELUDropoutTorchModel, self).__init__()
|
|
self.gelu = nn.GELU()
|
|
self.dropout = nn.Dropout(p=p)
|
|
|
|
def forward(self, x):
|
|
x = self.gelu(x)
|
|
x = self.dropout(x)
|
|
return x
|
|
|
|
def get_init_inputs():
|
|
"""获取模型初始化参数"""
|
|
return [0.1] # dropout概率
|
|
|
|
def get_inputs():
|
|
"""获取模型输入数据"""
|
|
torch.manual_seed(42)
|
|
return [torch.randn(128, 256)]
|
|
|
|
def test_gelu_dropout():
|
|
"""测试GELU+Dropout功能"""
|
|
print("=" * 60)
|
|
print("GELU+Dropout功能测试")
|
|
print("=" * 60)
|
|
|
|
# 创建测试数据
|
|
torch.manual_seed(42)
|
|
batch_size, hidden_size = 128, 256
|
|
input_tensor = torch.randn(batch_size, hidden_size)
|
|
|
|
# 创建模型
|
|
model = GELUDropoutTorchModel(dropout_prob=0.1)
|
|
|
|
# 训练模式测试
|
|
model.train()
|
|
train_output = model(input_tensor)
|
|
print(f"训练模式输出形状: {train_output.shape}")
|
|
print(f"训练模式输出范围: [{train_output.min().item():.3f}, {train_output.max().item():.3f}]")
|
|
|
|
# 推理模式测试
|
|
model.eval()
|
|
eval_output = model(input_tensor)
|
|
print(f"推理模式输出形状: {eval_output.shape}")
|
|
print(f"推理模式输出范围: [{eval_output.min().item():.3f}, {eval_output.max().item():.3f}]")
|
|
|
|
# 验证Dropout效果
|
|
zero_count_train = (train_output == 0).sum().item()
|
|
zero_count_eval = (eval_output == 0).sum().item()
|
|
|
|
print(f"训练模式零值比例: {zero_count_train / train_output.numel():.3f}")
|
|
print(f"推理模式零值比例: {zero_count_eval / eval_output.numel():.3f}")
|
|
|
|
# 验证GELU激活函数
|
|
gelu_only = 0.5 * input_tensor * (1 + torch.tanh(0.7978845608028654 * (input_tensor + 0.044715 * torch.pow(input_tensor, 3))))
|
|
print(f"纯GELU输出范围: [{gelu_only.min().item():.3f}, {gelu_only.max().item():.3f}]")
|
|
|
|
return train_output, eval_output
|
|
|
|
def main():
|
|
"""主测试函数"""
|
|
try:
|
|
train_output, eval_output = test_gelu_dropout()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("测试总结")
|
|
print("=" * 60)
|
|
print("✅ GELU+Dropout功能验证通过")
|
|
print("✅ 训练/推理模式切换正常")
|
|
print("✅ Dropout效果符合预期")
|
|
|
|
except Exception as e:
|
|
print(f"测试过程中出现错误: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |