Spaces:
Running
on
Zero
Running
on
Zero
# Copyright (c) Meta Platforms, Inc. and affiliates. | |
# All rights reserved. | |
# | |
# This source code is licensed under the license found in the | |
# LICENSE file in the root directory of this source tree. | |
import logging | |
from typing import Optional | |
import torch | |
import torch.nn as nn | |
from mmengine.runner import load_checkpoint | |
class AlexNet(nn.Module): | |
"""AlexNet backbone. | |
Args: | |
num_classes (int): number of classes for classification. | |
""" | |
def __init__(self, num_classes: int = -1): | |
super().__init__() | |
self.num_classes = num_classes | |
self.features = nn.Sequential( | |
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), | |
nn.ReLU(inplace=True), | |
nn.MaxPool2d(kernel_size=3, stride=2), | |
nn.Conv2d(64, 192, kernel_size=5, padding=2), | |
nn.ReLU(inplace=True), | |
nn.MaxPool2d(kernel_size=3, stride=2), | |
nn.Conv2d(192, 384, kernel_size=3, padding=1), | |
nn.ReLU(inplace=True), | |
nn.Conv2d(384, 256, kernel_size=3, padding=1), | |
nn.ReLU(inplace=True), | |
nn.Conv2d(256, 256, kernel_size=3, padding=1), | |
nn.ReLU(inplace=True), | |
nn.MaxPool2d(kernel_size=3, stride=2), | |
) | |
if self.num_classes > 0: | |
self.classifier = nn.Sequential( | |
nn.Dropout(), | |
nn.Linear(256 * 6 * 6, 4096), | |
nn.ReLU(inplace=True), | |
nn.Dropout(), | |
nn.Linear(4096, 4096), | |
nn.ReLU(inplace=True), | |
nn.Linear(4096, num_classes), | |
) | |
def init_weights(self, pretrained: Optional[str] = None) -> None: | |
if isinstance(pretrained, str): | |
logger = logging.getLogger() | |
load_checkpoint(self, pretrained, strict=False, logger=logger) | |
elif pretrained is None: | |
# use default initializer | |
pass | |
else: | |
raise TypeError('pretrained must be a str or None') | |
def forward(self, x: torch.Tensor) -> torch.Tensor: | |
x = self.features(x) | |
if self.num_classes > 0: | |
x = x.view(x.size(0), 256 * 6 * 6) | |
x = self.classifier(x) | |
return x | |