(
self,
in_channels: int,
embed_dim: int,
depth: int,
img_dim: int,
patch_size: int = 1,
has_text: bool = False,
num_classes=-1,
drop_path_rate=0.1,
n_context_token: int = 0,
d_context: int = 0,
ssm_cfg=None,
norm_epsilon: float = 1e-5,
rms_norm: bool = True,
fused_add_norm=True,
residual_in_fp32=True,
initializer_cfg=None,
scan_type="v2",
video_frames=0,
tpe=False, # apply temporal positional encoding for video-related task
device="cuda",
use_pe=0,
use_jit=True,
m_init=True,
use_checkpoint=False,
dtype=torch.float32,
)
| 547 | """ |
| 548 | |
| 549 | def __init__( |
| 550 | self, |
| 551 | in_channels: int, |
| 552 | embed_dim: int, |
| 553 | depth: int, |
| 554 | img_dim: int, |
| 555 | patch_size: int = 1, |
| 556 | has_text: bool = False, |
| 557 | num_classes=-1, |
| 558 | drop_path_rate=0.1, |
| 559 | n_context_token: int = 0, |
| 560 | d_context: int = 0, |
| 561 | ssm_cfg=None, |
| 562 | norm_epsilon: float = 1e-5, |
| 563 | rms_norm: bool = True, |
| 564 | fused_add_norm=True, |
| 565 | residual_in_fp32=True, |
| 566 | initializer_cfg=None, |
| 567 | scan_type="v2", |
| 568 | video_frames=0, |
| 569 | tpe=False, # apply temporal positional encoding for video-related task |
| 570 | device="cuda", |
| 571 | use_pe=0, |
| 572 | use_jit=True, |
| 573 | m_init=True, |
| 574 | use_checkpoint=False, |
| 575 | dtype=torch.float32, |
| 576 | ): |
| 577 | # assert num_classes == -1, "num_classes should be -1" |
| 578 | # assert n_context_token == 0, "n_context_token should be 0" |
| 579 | |
| 580 | self.factory_kwargs = factory_kwargs = {"device": device, "dtype": dtype} |
| 581 | super().__init__() |
| 582 | self.in_channels = in_channels |
| 583 | self.out_channels = in_channels |
| 584 | self.patch_size = patch_size |
| 585 | self.embed_dim = embed_dim |
| 586 | self.tpe = tpe |
| 587 | |
| 588 | self.residual_in_fp32 = residual_in_fp32 |
| 589 | self.fused_add_norm = fused_add_norm |
| 590 | self.video_frames = video_frames |
| 591 | self.use_pe = use_pe |
| 592 | num_patches = (img_dim // patch_size) ** 2 |
| 593 | self.use_checkpoint = use_checkpoint |
| 594 | print( |
| 595 | "use_checkpoint", |
| 596 | use_checkpoint, |
| 597 | "use_pe", |
| 598 | use_pe, |
| 599 | "use tpe", |
| 600 | tpe, |
| 601 | "num_patches", |
| 602 | num_patches, |
| 603 | "use_jit", |
| 604 | use_jit, |
| 605 | ) |
| 606 |
nothing calls this directly
no test coverage detected