Deep Learning
Pytorch how to use nn.ModuleDict with zip for iteration
ddokkddokk
2023. 7. 19. 10:15
반응형
class TEST(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(10, 10, 1)
def forward(self, x):
return self.conv(x)
num_res = 2
BEVEncoder = nn.ModuleDict()
UpSampler = nn.ModuleDict()
for _ in range(num_res):
IsSelfAttn = True if _ == 0 else False
BEVEncoder[str(_)] = TEST()
if (_ == 0):
UpSampler[str(_)] = None
else:
UpSampler[str(_)] = TEST()
for (_, enc), (_, up) in sorted(zip(BEVEncoder.items(), UpSampler.items()), reverse=True):
bp = 0