본문 바로가기

Deep Learning

Pytorch how to use nn.ModuleDict with zip for iteration

반응형
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