Hi All, I got this error when I do inference on CPU. It works well on GPU but when I use the CPU it throws this error.
This is the model architecture
class BasicBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride, activation):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=stride, # downsample with first conv
padding=1,
bias=False)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(
out_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.bn2 = nn.BatchNorm2d(out_channels)
self.activation = activation
self.shortcut = nn.Sequential()
if in_channels != out_channels:
self.shortcut.add_module(
'conv',
nn.Conv2d(
in_channels,
out_channels,
kernel_size=1,
stride=stride, # downsample
padding=0,
bias=False))
self.shortcut.add_module('bn', nn.BatchNorm2d(out_channels)) # BN
def forward(self, x):
y = self.activation(self.bn1(self.conv1(x)))
y = self.bn2(self.conv2(y))
y += self.shortcut(x)
y = self.activation(y) # apply ReLU after addition
return y
class DriverNet(nn.Module):
def __init__(self):
super(DriverNet, self).__init__()
self.activation = nn.ReLU()
self.conv_layers = nn.Sequential(
nn.BatchNorm2d(3),
BasicBlock(in_channels=3, out_channels=16, stride = 2, activation = self.activation),
BasicBlock(in_channels=16, out_channels=24, stride = 2, activation = self.activation),
BasicBlock(in_channels=24, out_channels=32, stride = 2, activation = self.activation),
BasicBlock(in_channels=32, out_channels=48, stride = 2, activation = self.activation),
BasicBlock(in_channels=48, out_channels=64, stride = 2, activation = self.activation),
BasicBlock(in_channels=64, out_channels=128, stride = 2, activation = self.activation),
# BasicBlock(in_channels=128, out_channels=256, stride = 2, activation = self.activation),
nn.Dropout(p=0.5)
)
self.linear_layers = nn.Sequential(
nn.Linear(in_features=128*4*5, out_features=128),
self.activation,
# nn.BatchNorm1d(128),
nn.Dropout(p=0.5),
nn.Linear(in_features=128, out_features=2),
)
def forward(self, x):
output = self.conv_layers(x)
output = output.view(output.size(0), -1)
output = self.linear_layers(output)
return output
This is the inference code:
device = torch.device("cpu") #torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = DriverNet().to(device)
model = torch.load("model.pt")
model = model.eval()
image = cv2.resize(image, dim, interpolation = cv2.INTER_AREA)
image_tensor = torch.from_numpy(image.transpose(2,0,1))
image_tensor = image_tensor.unsqueeze(0).float().to(device)
output = model(image_tensor)