forked from ccf-ai-infra/GPUCodeForces
87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
原始模型：InstanceNorm + Dropout
|
||
"""
|
||
def __init__(self, num_features=64, eps=1e-5, affine=True, dropout_p=0.1, track_running_stats=False):
|
||
super(Model, self).__init__()
|
||
self.num_features = num_features
|
||
self.eps = eps
|
||
self.affine = affine
|
||
self.dropout_p = dropout_p
|
||
self.track_running_stats = track_running_stats
|
||
|
||
# 创建InstanceNorm层
|
||
self.instance_norm = nn.InstanceNorm2d(
|
||
num_features=num_features,
|
||
eps=eps,
|
||
affine=affine,
|
||
track_running_stats=track_running_stats
|
||
)
|
||
|
||
# 创建Dropout层
|
||
self.dropout = nn.Dropout(p=dropout_p)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
原始实现：先InstanceNorm，再Dropout
|
||
|
||
Args:
|
||
x (torch.Tensor): 输入张量 shape [B, C, H, W]
|
||
|
||
Returns:
|
||
torch.Tensor: Dropout(InstanceNorm(x))
|
||
"""
|
||
x_normalized = self.instance_norm(x)
|
||
return self.dropout(x_normalized)
|
||
|
||
class ModelNew(torch.nn.Module):
|
||
"""
|
||
融合模型：直接实现InstanceNorm + Dropout
|
||
"""
|
||
def __init__(self, num_features=64, eps=1e-5, affine=True, dropout_p=0.1, track_running_stats=False):
|
||
super(ModelNew, self).__init__()
|
||
self.num_features = num_features
|
||
self.eps = eps
|
||
self.affine = affine
|
||
self.dropout_p = dropout_p
|
||
self.track_running_stats = False # 强制为False以支持CUDA实现
|
||
|
||
if affine:
|
||
self.weight = torch.nn.Parameter(torch.ones(num_features))
|
||
self.bias = torch.nn.Parameter(torch.zeros(num_features))
|
||
else:
|
||
self.register_parameter('weight', None)
|
||
self.register_parameter('bias', None)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
融合实现：在CUDA kernel中直接完成InstanceNorm + Dropout
|
||
|
||
Args:
|
||
x (torch.Tensor): 输入张量 shape [B, C, H, W]
|
||
|
||
Returns:
|
||
torch.Tensor: Dropout(InstanceNorm(x))
|
||
"""
|
||
# 这个将在CUDA中实现
|
||
pass
|
||
|
||
# 测试参数
|
||
batch_size = 128
|
||
num_features = 64
|
||
height = 128
|
||
width = 128
|
||
dropout_p = 0.1
|
||
|
||
def get_inputs():
|
||
"""生成测试输入"""
|
||
x = torch.randn(batch_size, num_features, height, width)
|
||
return [x]
|
||
|
||
def get_init_inputs():
|
||
"""获取初始化参数"""
|
||
return [num_features]
|