forked from ccf-ai-infra/GPUCodeForces
47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
GROUPS = 4
|
|
|
|
|
|
class ChannelShuffle(nn.Module):
|
|
|
|
def __init__(self, groups):
|
|
super().__init__()
|
|
self.groups = groups
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
batch_size, num_channels, height, width = x.size()
|
|
channels_per_group = num_channels // self.groups
|
|
|
|
# 1. Reshape
|
|
x = x.view(batch_size, self.groups, channels_per_group, height, width)
|
|
|
|
# 2. Transpose (交换 groups 和 channels_per_group 维度)
|
|
# dim 1 is groups, dim 2 is channels_per_group
|
|
x = torch.transpose(x, 1, 2).contiguous()
|
|
|
|
# 3. Flatten
|
|
x = x.view(batch_size, num_channels, height, width)
|
|
|
|
return x
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, groups=GROUPS):
|
|
super().__init__()
|
|
self.op = ChannelShuffle(groups)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.op(x)
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [GROUPS] |