nn.ModuleDict()

https://gaussian37.github.io/dl-pytorch-snippets/#block%EC%9D%84-%EC%8C%93%EA%B8%B0-%EC%9C%84%ED%95%9C-module-sequential-modulelist-moduledict-1

pytorch 기본 문법 및 코드, 팁 snippets

언제 사용하는가?

• 설정값들을 dictionary 형태로 선택하여 꺼내 써야 할 때

 

예시

def conv_block(in_f, out_f, activation='relu', *args, **kwargs):
    
    activations = nn.ModuleDict([
                ['lrelu', nn.LeakyReLU()],
                ['relu', nn.ReLU()]
    ])
    
    return nn.Sequential(
        nn.Conv2d(in_f, out_f, *args, **kwargs),
        nn.BatchNorm2d(out_f),
        activations[activation]
    )

print(conv_block(1, 32,'lrelu', kernel_size=3, padding=1))
print(conv_block(1, 32,'relu', kernel_size=3, padding=1))