forked from ccf-ai-infra/GPUCodeForces
41 lines
1021 B
Python
41 lines
1021 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
BATCH_SIZE = 8
|
|
CHANNELS = 16
|
|
DEPTH = 16 # D_in
|
|
HEIGHT = 16 # H_in
|
|
WIDTH = 16 # W_in
|
|
|
|
PADDING = (1, 1, 2, 2, 1, 0)
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self, padding):
|
|
super().__init__()
|
|
|
|
if isinstance(padding, int):
|
|
# F.pad 需要 6-tuple
|
|
self.padding_tuple = (padding,) * 6
|
|
else:
|
|
self.padding_tuple = padding
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# F.pad 5D 张量 (N, C, D, H, W)
|
|
# 填充顺序: (pad_W_L, pad_W_R, pad_H_T, pad_H_B, pad_D_F, pad_D_K)
|
|
# 这与 nn.ReflectionPad3d 的构造函数顺序一致
|
|
return F.pad(x, self.padding_tuple, mode='reflect')
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [PADDING]
|