forked from ccf-ai-infra/GPUCodeForces
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
BATCH_SIZE = 32
|
|
CHANNELS = 64
|
|
WIDTH = 128 # W_in
|
|
PADDING = (3, 1) # (padding_left, padding_right)
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
nn.ReflectionPad1d 的纯 PyTorch 基准实现
|
|
(使用 F.pad)
|
|
"""
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
# F.pad 需要 (left, right) 格式
|
|
self.padding_tuple = (padding, padding)
|
|
else:
|
|
self.padding_tuple = padding
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, pad_dim_1_left, ...)
|
|
# 因为我们只 pad 最后一个维度 (dim -1)，所以元组是 (pad_L, pad_R)
|
|
return F.pad(x, self.padding_tuple, mode='reflect')
|
|
|
|
|
|
def get_inputs():
|
|
"""
|
|
生成一个 (N, C, W) 形状的输入
|
|
"""
|
|
x = torch.randn(BATCH_SIZE, CHANNELS, WIDTH, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [PADDING]
|