def pre_fcn_resnet101(in_channel, out_channel):
model = fcn_resnet101(pretrained=False, progress=False)
url = "https://download.pytorch.org/models/fcn_resnet101_coco-7ecb50ca.pth" # COCO
model_dict = model.state_dict()
pretrained_dict = model_zoo.load_url(url, progress=False)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model.backbone.conv1 = nn.Conv2d(in_channel,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
model.classifier = FCNHead(2048, out_channel)
return model
def test_fcn(self):
x = Variable(torch.randn(BATCH_SIZE, 3, 224, 224).fill_(1.0))
self.exportTest(toC(fcn_resnet101()), toC(x), rtol=1e-3, atol=1e-5)