This is my code:
import torch
import torch.nn as nn
class AlexNet(nn.Module):
def __init__(self, __output_size):
super(AlexNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=96, kernel_size=11, stride=4),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2)
)
self.layer2 = nn.Sequential(
nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2)
)
self.layer3 = nn.Sequential(
nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True)
)
self.layer4 = nn.Sequential(
nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True)
)
self.layer5 = nn.Sequential(
nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=1)
)
self.layer6 = nn.Sequential(
nn.Dropout(p=0.5)
)
#self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
self.layer7 = nn.Sequential(
nn.Linear(in_features=4096, out_features=4096, bias=True),
nn.ReLU()
)
self.layer8 = nn.Sequential(
nn.Dropout(p=0.5)
)
self.layer9 = nn.Sequential(
nn.Linear(in_features=4096, out_features=4096, bias=True),
nn.ReLU()
)
self.layer10 = nn.Sequential(
nn.Linear(in_features=4096, out_features=1000, bias=True),
nn.Softmax()
)
def forward(self, x):
_output = self.layer1(x)
_output = self.layer2(_output)
_output = self.layer3(_output)
_output = self.layer4(_output)
_output = self.layer5(_output)
_output = self.layer6(_output)
# _output = self.avgpool(_output)
_output = torch.flatten(_output, 1)
_output = self.layer7(_output)
_output = self.layer8(_output)
_output = self.layer9(_output)
_output = self.layer10(_output)
return _output
This is my training part:
def train(_train_loader, _model, _num_epochs, _device, _criterion, _optimizer):
try:
_total_steps = len(_train_loader)
for _epochs in range(_num_epochs):
for i, (_images, _labels) in enumerate(_train_loader):
_images = _images.to(_device)
_labels = _labels.to(_device)
# forward pass
_outputs = _model(_images)
_loss = _criterion(_outputs, _labels)
# backward pass and optimization
_optimizer.zero_grad()
_loss.backward()
_optimizer.step()
if (i+1) % 100 == 0:
print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(_epochs+1, _num_epochs, i+1, _total_steps, _loss.item()))
except Exception as error:
print("An error occured while training")
print(error)
raise error
Initially I was getting:
RuntimeError: Given groups=1, weight of size [96, 3, 11, 11], expected input[100, 1, 28, 28] to have 3 channels, but got 1 channels instead
Then I changed the number of channels to 1, I got the above mentioned error. Please help me out to resolve this error!