def main(argv):
weight_loader = load_weights.WeightLoader(mapping=dummy_mapping)
if FLAGS.test_print_tensors:
logging.info("Testing print tensor...")
weight_loader.print_tensors()
logging.info("Print tensor done!")
if FLAGS.test_map_tensors:
logging.info("Testing load variable...")
weight_loader.load_variable("test")
logging.info("Testing variable done!")
if FLAGS.test_mapping:
logging.info("Testing mapping...")
model = models.resnet50()
mapping = default_mapping.get_default_mapping()
weight_loader = load_weights.WeightLoader(mapping=mapping)
weight_loader.validate_mapping(model)
logging.info("Testing mapping done!")
if FLAGS.test_model:
logging.info("Testing model...")
model = models.resnet50()
mapping = default_mapping.get_default_mapping()
weight_loader = load_weights.WeightLoader(mapping=mapping)
model = weight_loader.set_torch_model(model)
logging.info("Testing model done!")
def __init__(self, encoded_channels, input_channels, no_upsampling=False, load_tf_weights=True, **resnet_args):
super().__init__()
self.resnet = torchvision_resnet.resnet50(
pretrained=True, **resnet_args).eval().requires_grad_(False)
self.no_upsampling = no_upsampling
if load_tf_weights:
mapping = default_mapping.get_default_mapping()
weight_loader = load_weights.WeightLoader(mapping=mapping)
self.resnet = weight_loader.set_torch_model(self.resnet)
self.encoded_channels = encoded_channels
self.input_channels = input_channels
self.stages = ('early', 'middle', 'deep')
self.resnet_extractor = ResnetActivationExtractor(
self.resnet, conv_layer=None)
self.resnet_correspondence_extractor = ResnetCorrespondenceExtractor(
self.resnet_extractor)
self.encoder, self.decoder = self.get_encoder_decoder()
# we need this such that the encoders get tranfered to gpu automatically
self.e1 = self.encoder['early']
self.e2 = self.encoder['middle']
self.e3 = self.encoder['deep']
self.d1 = self.decoder['early']
self.d2 = self.decoder['middle']
self.d3 = self.decoder['deep']
def __init__(self, load_tf_weights=True):
super().__init__()
self.resnet = torchvision.models.resnet50(
pretrained=True).eval().requires_grad_(False)
if load_tf_weights:
mapping = default_mapping.get_default_mapping()
weight_loader = load_weights.WeightLoader(mapping=mapping)
self.resnet = weight_loader.set_torch_model(self.resnet)
self.encoded_channels = 64
self.stages = ('early', 'middle', 'deep')
self.input_channels = {
'early': 512,
'middle': 1024,
'deep': 2048,
}
self.resnet_extractor = ResnetActivationExtractor(self.resnet,
conv_layer=None)
self.resnet_correspondence_extractor = ResnetCorrespondenceExtractor(
self.resnet_extractor)
self.encoder = {
'early':
nn.Sequential(
nn.Conv2d(in_channels=self.input_channels['early'],
out_channels=self.encoded_channels,
kernel_size=(1, 1)),
nn.ReLU(True),
),
'middle':
nn.Sequential(
nn.Conv2d(in_channels=self.input_channels['middle'],
out_channels=self.encoded_channels,
kernel_size=(1, 1)),
nn.ReLU(True),
),
'deep':
nn.Sequential(
nn.Conv2d(in_channels=self.input_channels['deep'],
out_channels=self.encoded_channels,
kernel_size=(1, 1)),
nn.ReLU(True),
),
}
self.decoder = {
'early':
nn.Sequential(
nn.ConvTranspose2d(in_channels=self.encoded_channels,
out_channels=self.input_channels['early'],
kernel_size=(1, 1)),
nn.BatchNorm2d(self.input_channels['early']),
nn.ReLU(True),
),
'middle':
nn.Sequential(
nn.ConvTranspose2d(in_channels=self.encoded_channels,
out_channels=self.input_channels['middle'],
kernel_size=(1, 1)),
nn.BatchNorm2d(self.input_channels['middle']),
nn.ReLU(True),
),
'deep':
nn.Sequential(
nn.ConvTranspose2d(in_channels=self.encoded_channels,
out_channels=self.input_channels['deep'],
kernel_size=(1, 1)),
nn.BatchNorm2d(self.input_channels['deep']),
nn.ReLU(True),
),
}
# we need this such that the encoders get tranfered to gpu automatically
self.e1 = self.encoder['early']
self.e2 = self.encoder['middle']
self.e3 = self.encoder['deep']
self.d1 = self.decoder['early']
self.d2 = self.decoder['middle']
self.d3 = self.decoder['deep']