forked from ccf-ai-infra/GPUCodeForces
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
BATCH_SIZE = 32
|
|
CHANNELS = 64
|
|
HEIGHT = 32 # H_in
|
|
WIDTH = 32 # W_in
|
|
|
|
# (pad_L, pad_R, pad_T, pad_B)
|
|
PADDING = (1, 1, 2, 0)
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
nn.ReflectionPad2d 的纯 PyTorch 基准实现
|
|
(使用 F.pad)
|
|
"""
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
# F.pad 需要 (left, right, top, bottom) 格式
|
|
self.padding_tuple = (padding, padding, 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, ...)
|
|
# 对应 (N, C, H, W)，我们需要 pad 最后两个维度
|
|
# F.pad 接受的顺序是 (pad_W_left, pad_W_right, pad_H_top, pad_H_bottom)
|
|
return F.pad(x, self.padding_tuple, mode='reflect')
|
|
|
|
|
|
def get_inputs():
|
|
"""
|
|
生成一个 (N, C, H, W) 形状的输入
|
|
"""
|
|
x = torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [PADDING]
|