def __init__(self, num_classes=1000):
super(Resnet18_8s, self).__init__()
# Load the pretrained weights, remove avg pool
# layer and get the output stride of 8
resnet18_8s = models.resnet18(fully_conv=True,
pretrained=True,
output_stride=8,
remove_avg_pool_layer=True)
# Randomly initialize the 1x1 Conv scoring layer
resnet18_8s.fc = nn.Conv2d(resnet18_8s.inplanes, num_classes, 1)
self.resnet18_8s = resnet18_8s
self._normal_initialization(self.resnet18_8s.fc)
def get_resnet18(device: torch.device) -> GetterReturnType:
N = 32
model = models.resnet18(pretrained=False)
criterion = torch.nn.CrossEntropyLoss()
model.to(device)
params, names = extract_weights(model)
inputs = torch.rand([N, 3, 224, 224], device=device)
labels = torch.rand(N, device=device).mul(10).long()
def forward(*new_params: Tensor) -> Tensor:
load_weights(model, names, new_params)
out = model(inputs)
loss = criterion(out, labels)
return loss
return forward, params
def get_resnet18(device: torch.device) -> GetterReturnType:
N = 32
model = models.resnet18(pretrained=False)
if has_functorch:
from functorch.experimental import replace_all_batch_norm_modules_
replace_all_batch_norm_modules_(model)
criterion = torch.nn.CrossEntropyLoss()
model.to(device)
params, names = extract_weights(model)
inputs = torch.rand([N, 3, 224, 224], device=device)
labels = torch.rand(N, device=device).mul(10).long()
def forward(*new_params: Tensor) -> Tensor:
load_weights(model, names, new_params)
out = model(inputs)
loss = criterion(out, labels)
return loss
return forward, params