Can we use Neural ODE Block with a normal neural network block?

[suyashsachdeva]

Hi, I was trying to make a model in which I want to use a convolutional neural ODE block followed by a convolutional block and I am unable to understand what is going wrong.

[massastrello]

Hi! Would you share a minimal example to reproduce your issue? (please paste the actual code and not the screenshots)

[suyashsachdeva]

class UPBlock(nn.Module):

def __init__(self, infilter:int, outfilter:int, kernel:int, moment:float):
    super(UPBlock, self).__init__()
    self.conv1 = nn.Conv2d(infilter, outfilter, kernel)
    self.norm1 = nn.BatchNorm2d(outfilter, momentum=moment)

    self.conv2 = nn.Conv2d(outfilter, outfilter, kernel)
    self.norm2 = nn.BatchNorm2d(outfilter, momentum=moment)

    self.act = nn.ReLU()
    self.pad = nn.ZeroPad2d(int((kernel-1)/2.0))
    self.pool = nn.MaxPool2d(2, 2)

def forward(self, x):
    x = self.act(self.norm1(self.pad(self.conv1(x))))
    x = self.act(self.norm2(self.pad(self.conv2(x))))
    return self.pool(x)

class DenseBlock(nn.Module):

def __init__(self, filter:int, dense:int, drop:float, classes:int):
    super(DenseBlock, self).__init__()

    self.pool = nn.AdaptiveAvgPool2d((1,1))
    self.flat = nn.Flatten()
    self.dense = nn.Linear(filter, dense)
    self.drop = nn.Dropout(drop)
    self.final = nn.Linear(dense, classes)

def forward(self,x):
    x = self.dense(self.drop(self.flat(self.pool(x))))
    return F.softmax(self.final(self.drop(F.relu(x))))

    self.final = DenseBlock(filter*pow(gf, num-1), dense, drop, classes)

class CNNODE(nn.Module):

def __init__(self, num=3, filter=64, dense=256, classes=15, gf=2, kernel=3, moment=0.9, drop=0.2):
    super(CNNODE, self).__init__()

    self.cnnode = nn.ModuleList([])
    self.upsamp = nn.ModuleList([])
    self.final = DenseBlock(filter*pow(gf, num-1), dense, drop, classes)
    self.conv = UPBlock(1, filter, kernel, moment)

    for _ in range(num):
        f = self.func(filter, filter*2, kernel, moment)
        model = NeuralODE(f, sensitivity='adjoint', solver='rk4', solver_adjoint='dopri5', atol_adjoint=1e-4, rtol_adjoint=1e-4)
        self.cnnode.append(model)
        self.upsamp.append(UPBlock(filter, filter*2, kernel, moment))
        filter = filter*gf

def forward(self, x, t_span):
    x = self.conv(x)
    for neuralode, neuralnetwork in zip(self.cnnode, self.upsamp):
        x = neuralode(x, t_span)
        x = neuralnetwork()
    return self.final(x)  

def func(self, infilter, outfilter, kernel, moment):
    return nn.Sequential(nn.Conv2d(infilter, outfilter, kernel), 
                        nn.ZeroPad2d(int((kernel-1)/2)),
                        nn.BatchNorm2d(outfilter, momentum=moment),
                        nn.ReLU(), 
                        nn.Conv2d(outfilter, infilter, kernel),
                        nn.ZeroPad2d(int((kernel-1)/2)),
                        nn.BatchNorm2d(infilter, momentum=moment),
                        nn.ReLU(), 
            )    

class Learner(LightningModule):

def __init__(self, t_span:th.tensor, model:nn.Module):
    super(Learner, self).__init__()

    self.model, self.t_span = model, t_span

def forward(self, x):
    return self.model(x)

def training_step(self, batch, batch_idx):
    x, y = batch
    t_eval, y_hat = self.model(x, self.t_span)
    y_hat = y_hat[-1] # select last point of solution trajectory
    loss = nn.CrossEntropyLoss()(y_hat, y)
    return {'loss': loss}

def configure_optimizers(self):
    return th.optim.Adam(self.model.parameters(), lr=0.0001)

def train_dataloader(self):
    return trainloader

model = CNNODE() t_span = th.linspace(0, 1, 100) learn = Learner(t_span, model)

trainer = pl.Trainer(min_epochs=200, max_epochs=300) trainer.fit(learn, train_dataloaders=CircuitDataset)

TypeError Traceback (most recent call last) ~\AppData\Local\Temp/ipykernel_3100/2893833635.py in 1 trainer = pl.Trainer(min_epochs=200, max_epochs=300) ----> 2 trainer.fit(learn, train_dataloaders=CircuitDataset)

c:\Users\suyash\AppData\Local\Programs\Python\Python39\lib\site-packages\pytorch_lightning\trainer\trainer.py in fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path) 694 """ 695 self.strategy.model = model --> 696 self._call_and_handle_interrupt( 697 self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path 698 )

c:\Users\suyash\AppData\Local\Programs\Python\Python39\lib\site-packages\pytorch_lightning\trainer\trainer.py in _call_and_handle_interrupt(self, trainer_fn, *args, kwargs) 648 return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, *kwargs) 649 else: --> 650 return trainer_fn(args, kwargs) 651 # TODO(awaelchli): Unify both exceptions below, where KeyboardError doesn't re-raise 652 except KeyboardInterrupt as exception:

c:\Users\suyash\AppData\Local\Programs\Python\Python39\lib\site-packages\pytorch_lightning\trainer\trainer.py in _fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path) 733 ckpt_path, model_provided=True, model_connected=self.lightning_module is not None 734 ) --> 735 results = self._run(model, ckpt_path=self.ckpt_path) 736 ...

• (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, tuple of ints padding, tuple of ints dilation, int groups) didn't match because some of the arguments have invalid types: (!str!, !Parameter!, !Parameter!, !tuple!, !tuple!, !tuple!, int)
• (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, str padding, tuple of ints dilation, int groups) didn't match because some of the arguments have invalid types: (!str!, !Parameter!, !Parameter!, !tuple!, !tuple!, !tuple!, int)

[suyashsachdeva]

Thank you I found the mistake