class VisionTransformer(nn.Module):
""" Vision Transformer
A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
- https://arxiv.org/abs/2010.11929
Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
- https://arxiv.org/abs/2012.12877
"""
def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, distilled=False,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None,
act_layer=None, weight_init=''):
super().__init__()
self.num_classes = num_classes
self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
self.num_tokens = 2 if distilled else 1
norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
act_layer = act_layer or nn.GELU
self.register_buffer('matrix', torch.zeros(12, 197, 192)) #이부분입니다.
vision transformer 코드인데, 위 코드에서 self.register_buffer('matrix', torch.zeros(12, 197, 192))로 선언할 때와 self.matrix = torch.zeros(12, 197, 192).cuda() 로 선언할때와 성능 차이가 발생합니다. checkpoint에 값을 저장하기 위해서 register_buffer를 사용하고자하는데, register_buffer 사용 안할 때 성능이 더 좋게나오더라구요.
위 matrix를 연산하는 곳은 한 줄 밖에 없는데, with torch.no_grad():로 연산을 수행합니다.
gradient가 흐르는건 아닌거같은데, 어떤 차이때문에 성능 차이가 발생하나요? 도저히 이해가 안됩니다.....
ddp를 사용중인데 이것에 문제가 존재하는 걸까요?