rjbruin / flexconv Goto Github PK

import torch
import torch.nn as nn
import torch.nn.functional as F
import ckconv
from ckconv.nn import CKConv
from omegaconf import OmegaConf


kernel_config = OmegaConf.create({"type": "MLP", "dim_linear": 2, "no_hidden": 2, "no_layers": 3, "activ_function": "ReLU","norm": "BatchNorm","omega_0": 1,"learn_omega_0": False,"weight_norm": False,"steerable": False,"init_spatial_value": 1.0,"bias_init": None,"input_scale": 25.6,"sampling_rate_norm": 1.0,"regularize": False,"regularize_params": {"res": 0 ,"res_offset": 0,"target": "gabor+mask","fn": "l2_relu","method":"together","factor": 0.001,"gauss_stddevs": 2.0,"gauss_factor": 0.5},"srf": {"scale": 0.}})


conv_config = OmegaConf.create({"type": "","use_fft": False, "bias": True,"padding": "same","stride": 1,"horizon": "same","cache": False })

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.conv1 = CKConv(3, 6, kernel_config, conv_config) # nn.Conv2d(3, 6, 5) --> original conv that works
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        print("x: ", x.shape)
        y = self.conv1(x)
        print("y: ", y.shape)
        x = self.pool(F.relu(y))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()


inn = torch.randn((1,3, 28, 28))
out = net(inn)