i am getting this error at the code below duing the training , the error in the model since the shapes as following
image.shape = torch.Size([2, 1, 512, 512])
masks.shape = torch.Size([2, 1, 512, 512])
outputs.shape =torch.Size([2, 1, 324, 324])
def train_model(model, data_train, criterion, optimizer):
"""Train the model and report validation error with training error
Args:
model: the model to be trained
criterion: loss function
data_train (DataLoader): training dataset
"""
model.train()
for batch, (images, masks) in enumerate(data_train):
images = Variable(images.float().cuda())
print(images.shape)
masks = Variable(masks.float().cuda())
print(masks.shape)
outputs = model(images)
print(outputs.shape)
#print(masks.shape, outputs.shape)
loss = criterion(outputs, masks).float()
optimizer.zero_grad()
loss.backward()
# Update weights
optimizer.step()
and here is my model SegNet , what is the worng in the model that makes the 'outputs ’ in shape 324?
class SegNet(nn.Module):
def __init__(self):
super(SegNet, self).__init__()
self.encoder_1 = nn.Sequential(
nn.Conv2d(3, 64, 7, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d((2, 2), stride=(2, 2), return_indices=True)
) # first group
self.encoder_2 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d((2, 2), stride=(2, 2), return_indices=True)
) # second group
self.encoder_3 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d((2, 2), stride=(2, 2), return_indices=True)
) # third group
self.encoder_4 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d((2, 2), stride=(2, 2), return_indices=True)
) # fourth group
self.unpool_1 = nn.MaxUnpool2d(2, stride=2) # get masks
self.unpool_2 = nn.MaxUnpool2d(2, stride=2)
self.unpool_3 = nn.MaxUnpool2d(2, stride=2)
self.unpool_4 = nn.MaxUnpool2d(2, stride=2)
self.decoder_1 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64)
) # first group
self.decoder_2 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64)
) # second group
self.decoder_3 = nn.Sequential(
nn.Conv2d(64, 64, 7, padding=3),
nn.BatchNorm2d(64)
) # third group
self.decoder_4 = nn.Sequential(
nn.Conv2d(64, 1, 7, padding=3),
nn.BatchNorm2d(1)
) # fourth group
def weight_init(self):
for i,m in enumerate(self.modules()):
if isinstance(m,nn.Conv2d):
nn.init.xavier_normal(m.weight)
nn.init.constant(m.bias, 0)
def forward(self, x):
size_1 = x.size()
x, indices_1 = self.encoder_1(x)
size_2 = x.size()
x, indices_2 = self.encoder_2(x)
size_3 = x.size()
x, indices_3 = self.encoder_3(x)
size_4 = x.size()
x, indices_4 = self.encoder_4(x)
x = self.unpool_1(x, indices_4, output_size=size_4)
x = self.decoder_1(x)
x = self.unpool_2(x, indices_3, output_size=size_3)
x = self.decoder_2(x)
x = self.unpool_3(x, indices_2, output_size=size_2)
x = self.decoder_3(x)
x = self.unpool_4(x, indices_1, output_size=size_1)
x = self.decoder_4(x)
x = F.sigmoid(x)
return x