| import torch |
| import torch.nn as nn |
| import numpy as np |
|
|
| class HyperspectralPatchEmbed(nn.Module): |
| """Image to Patch Embedding.""" |
|
|
| def __init__( |
| self, |
| patch_size: int = 16, |
| embed_dim: int = 768, |
| ): |
| super().__init__() |
| self.patch_size = patch_size |
| |
| |
| |
| |
| |
| |
| |
| |
| self.proj = nn.Sequential( |
| nn.Conv3d( |
| 1, |
| embed_dim*2, |
| kernel_size=(1, patch_size, patch_size), |
| stride=(1, patch_size, patch_size), |
| ), |
| nn.ReLU(), |
| nn.Conv3d( |
| embed_dim*2, |
| embed_dim, |
| kernel_size=(1, 1, 1), |
| stride=(1, 1, 1), |
| ), |
| ) |
|
|
| self.initialize_weights() |
|
|
| def initialize_weights(self): |
| try: |
| w = self.proj.weight.data |
| torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1])) |
| except: |
| w1 = self.proj[0].weight.data |
| w2 = self.proj[2].weight.data |
| torch.nn.init.xavier_uniform_(w1.view([w1.shape[0], -1])) |
| torch.nn.init.xavier_uniform_(w2.view([w2.shape[0], -1])) |
|
|
| def forward(self, x): |
| |
| assert x.shape[2] % self.patch_size == 0, "image height should be divisible by patch size" |
| assert x.shape[3] % self.patch_size == 0, "image width should be divisible by patch size" |
| x = self.proj(x.unsqueeze(1)) |
| B, D, C, H, W = x.shape |
| return x.permute(0, 2, 3, 4, 1).reshape(B, C, H*W, D) |
|
|
