glefundes / mobile-face-gaze

Lightweight gaze estimation with PyTorch.
GNU General Public License v2.0
94 stars 19 forks source link

AttributeError: module 'torchvision.models.mobilenet' has no attribute 'ConvBNReLU' #5

Open wangji9 opened 2 years ago

thonglv21 commented 2 years ago

conda install fastai pytorch=1.4.0 -c fastai -c pytorch -c conda-forge Worked for me.

prajotsl123 commented 2 years ago

Do this change

class ConvBNReLU(nn.Sequential):
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
        padding = (kernel_size - 1) // 2
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
            norm_layer(out_planes),
            nn.ReLU6(inplace=True)
        )

class GazeNet(nn.Module):

    def __init__(self, device):    
        super(GazeNet, self).__init__()
        self.device = device
        self.preprocess = transforms.Compose([
            transforms.Resize((112,112)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        model = models.mobilenet_v2(pretrained=True)
        model.features[-1] = ConvBNReLU(320, 256, kernel_size=1)

This will work