When I want to train my model, output tensor will be always zero. What am I doing wrong?
class CFD_CNN(nn.Module):
def __init__(self, out_ch=300):
super(CFD_CNN, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, out_ch, kernel_size=5, stride=5, padding=0),
nn.ReLU(True),
nn.Conv2d(out_ch, out_ch, kernel_size=5, stride=5, padding=0),
nn.ReLU(True),
nn.Conv2d(out_ch, out_ch, kernel_size=2, stride=2, padding=0),
nn.ReLU(True),
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(out_ch, out_ch, kernel_size=2, stride=2, padding=0),
nn.ReLU(True),
nn.ConvTranspose2d(out_ch, out_ch, kernel_size=5, stride=5, padding=0),
nn.ReLU(True),
nn.ConvTranspose2d(out_ch, 1, kernel_size=5, stride=5, padding=0),
nn.ReLU(True)
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
``
##############################################################
def train(epochs):
net.train()
for epoch in range(epochs):
for i, (distfunc, flowfields) in enumerate(train_loader):
distfunc = distfunc.unsqueeze(dim=1)
distfunc = distfunc.to(device=device, dtype = torch.float)
flowfields = flowfields.unsqueeze(dim=1)
flowfields = flowfields.to(device=device, dtype = torch.float)
output = net(distfunc)
print(output)
print(flowfields)
optimizer.zero_grad()
loss = loss_func(output, flowfields)
loss.backward()
optimizer.step()
print("Epoch: {}/{} | Batch: {}/{} | Loss: {}"
.format(epoch+1, epochs, i, len(train_loader), loss))