mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
Modified ViT get_intermediate_layers() to support dynamic image size
This commit is contained in:
parent
6e6f3686a7
commit
4731e4efc4
@ -667,9 +667,12 @@ class VisionTransformer(nn.Module):
|
|||||||
outputs = [out[:, self.num_prefix_tokens:] for out in outputs]
|
outputs = [out[:, self.num_prefix_tokens:] for out in outputs]
|
||||||
|
|
||||||
if reshape:
|
if reshape:
|
||||||
grid_size = self.patch_embed.grid_size
|
patch_size = self.patch_embed.patch_size
|
||||||
|
batch, _, height, width = x.size()
|
||||||
outputs = [
|
outputs = [
|
||||||
out.reshape(x.shape[0], grid_size[0], grid_size[1], -1).permute(0, 3, 1, 2).contiguous()
|
out.reshape(batch, int(math.ceil(height / patch_size[0])), int(math.ceil(width / patch_size[1])), -1)
|
||||||
|
.permute(0, 3, 1, 2)
|
||||||
|
.contiguous()
|
||||||
for out in outputs
|
for out in outputs
|
||||||
]
|
]
|
||||||
|
|
||||||
|
Loading…
x
Reference in New Issue
Block a user