MCPcopy Create free account
hub / github.com/FreeformRobotics/OTS / __init__

Method __init__

models/ots.py:16–26  ·  view source on GitHub ↗
(self, in_dim, compress)

Source from the content-addressed store, hash-verified

14
15class Obj_Attn_Block(Module):
16 def __init__(self, in_dim, compress):
17 super(Obj_Attn_Block, self).__init__()
18 channel_in = in_dim//int(2*compress)
19 self.value_conv = Conv2d(in_channels=in_dim, out_channels=channel_in, kernel_size=1)
20 self.query_conv = Conv2d(in_channels=channel_in, out_channels=channel_in, kernel_size=1)
21 self.key_conv = Conv2d(in_channels=channel_in, out_channels=channel_in, kernel_size=1)
22 self.gamma = Parameter(torch.zeros(1), requires_grad=True)
23 self.softmax = Softmax(dim=-1)
24
25 for layer in [self.value_conv, self.query_conv, self.key_conv]:
26 weight_init(layer)
27
28 def forward(self, x):
29 m_batchsize, C, length, _ = x.size()

Callers

nothing calls this directly

Calls 2

weight_initFunction · 0.85
__init__Method · 0.45

Tested by

no test coverage detected