How to use the torchx.layers.Conv2d function in torchx

To help you get started, we’ve selected a few torchx examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github SurrealAI / surreal / surreal / model / model_builders / builders.py View on Github external
def __init__(self, D_obs, D_out, conv_channels=[16, 32], kernel_sizes=[8, 4], strides=[4,2]):
        super(CNNStemNetwork, self).__init__()
        layers = []
        for i in range(len(conv_channels)):
            layers.append(L.Conv2d(conv_channels[i], kernel_size=kernel_sizes[i], stride=strides[i]))
            layers.append(L.ReLU())
        layers.append(L.Flatten())
        layers.append(L.Linear(D_out))
        layers.append(L.ReLU())
        self.model = L.Sequential(*layers)

        # instantiate parameters
        self.model.build((None, *D_obs))