class Mlp(nn.Module): """ Multilayer perceptron.""" def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features) self.act = act_layer() self.fc2 = nn.Linear(hidden_features, out_features) self.drop = nn.Dropout(drop) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.drop(x) x = self.fc2(x) x = self.drop(x) return x def window_partition(x, window_size): """ Args: x: (B, D, H, W, C) window_size (tuple[int]): window size Returns: windows: (B*num_windows, window_size*window_size, C) """ B, D, H, W, C = x.shape x = x.view(B, D // window_size[0], window_size[0], H // window_size[1], window_size[1], W // window_size[2], window_size[2], C) windows = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, reduce(mul, window_size), C) return windows def window_reverse(windows, window_size, B, D, H, W): """ Args: windows: (B*num_windows, window_size, window_size, C) window_size (tuple[int]): Window size H (int): Height of image W (int): Width of image Returns: x: (B, D, H, W, C) """ x = windows.view(B, D // window_size[0], H // window_size[1], W // window_size[2], window_size[0], window_size[1], window_size[2], -1) x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1) return x def get_window_size(x_size, window_size, shift_size=None): use_window_size = list(window_size) if shift_size is not None: use_shift_size = list(shift_size) for i in range(len(x_size)): if x_size[i] <= window_size[i]: use_window_size[i] = x_size[i] if shift_size is not None: use_shift_size[i] = 0 if shift_size is None: return tuple(use_window_size) else: return tuple(use_window_size), tuple(use_shift_size)
时间: 2023-05-21 20:02:16 浏览: 212
神经网络_Fashion_mnist_alexnet_CNN_
这是一个关于 PyTorch 的代码,其中包含了一个 MLP 模型和一些用于窗口分割和反转的函数。其中 MLP 模型包含了一个线性层和一个激活函数,以及一个 dropout 层。窗口分割和反转的函数用于将输入张量分割成多个窗口,并将这些窗口重新组合成一个张量。这些函数可以用于实现一些图像处理任务,例如局部感知机制。
阅读全文