forked from ccf-ai-infra/GPUCodeForces
31 lines
871 B
Python
31 lines
871 B
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
class Model(torch.nn.Module):
|
|
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, input_image: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
|
|
return F.grid_sample(input_image, grid, mode='bilinear', padding_mode='zeros', align_corners=True)
|
|
|
|
def get_inputs():
|
|
# 创建测试数据
|
|
N, C, H_in, W_in = 1, 3, 256, 256
|
|
H_out, W_out = 512, 512
|
|
|
|
input_image = torch.randn(N, C, H_in, W_in)
|
|
|
|
grid_y, grid_x = torch.meshgrid(
|
|
torch.linspace(-1, 1, H_out),
|
|
torch.linspace(-1, 1, W_out),
|
|
indexing='ij'
|
|
)
|
|
grid = torch.stack((grid_x, grid_y), dim=-1) # Shape: (H_out, W_out, 2)
|
|
grid = grid.unsqueeze(0).repeat(N, 1, 1, 1) # Shape: (N, H_out, W_out, 2)
|
|
|
|
return [input_image, grid]
|
|
|
|
def get_init_inputs():
|
|
return []
|