def make(args):
nonlinearity, kwargs = nn.LeakyReLU, {"inplace": True}
encoder_core = vgg.make_layers(vgg.vgg_cfg[args.model_type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
encoder = knn.Unit(args.model_in_channels, args.model_z_channels,
encoder_core)
decoder_core = vgg.make_layers(vgg.decoder_cfg[args.model_type])
decoder = knn.Unit(args.model_z_channels + args.model_keypoints,
args.model_in_channels, decoder_core)
keypoint_core = vgg.make_layers(vgg.vgg_cfg[args.model_type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
keypoint = knn.Unit(args.model_in_channels, args.model_keypoints,
keypoint_core)
keymapper = knn.GaussianLike(sigma=0.1)
kp_network = KeyNet(encoder,
keypoint,
keymapper,
decoder,
init_weights=True)
if args.load is not None:
kp_network.load(args.load)
if args.transfer_load is not None:
kp_network.load_from_autoencoder(args.transfer_load)
return kp_network
def make(type,
in_channels,
z_channels,
keypoints,
combine_mode='max',
load=None,
transfer_load=None,
map_device=None,
sigma=0.1):
"""
:param type
:param in_channels
:param channels
:param keypoints
:param conbine_mode: max, sum_and_clamp, loop, pretrained-network
:param map_device:
:return:
"""
nonlinearity, kwargs = nn.LeakyReLU, {"inplace": True}
encoder_core = vgg.make_layers(vgg.vgg_cfg[type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
encoder = knn.Unit(in_channels, z_channels, encoder_core)
decoder_core = vgg.make_layers(vgg.decoder_cfg[type])
decoder = knn.Unit(z_channels, in_channels, decoder_core)
keypoint_core = vgg.make_layers(vgg.vgg_cfg[type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
keypoint = knn.Unit(in_channels, keypoints, keypoint_core)
keymapper = knn.GaussianLike(sigma=sigma)
transporter_net = TransporterNet(encoder,
keypoint,
keymapper,
decoder,
init_weights=True,
combine_method=combine_mode)
if load is not None:
transporter_net.load(load, map_device)
if transfer_load is not None:
transporter_net.load_from_autoencoder(transfer_load)
return transporter_net
""" adds 2 channels that carry co-ordinate information """
b, h, w = x.size(0), x.size(2), x.size(3)
hm = torch.linspace(0, 1, h, dtype=x.dtype,
device=x.device).reshape(1, 1, h,
1).repeat(b, 1, 1, w)
wm = torch.linspace(0, 1, w, dtype=x.dtype,
device=x.device).reshape(1, 1, 1,
w).repeat(b, 1, h, 1)
return torch.cat((x, hm, wm), dim=1)
""" model """
nonlinearity, kwargs = nn.LeakyReLU, {"inplace": True}
encoder_core = vgg.make_layers(vgg.vgg_cfg[args.model_type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
encoder = knn.Unit(args.model_in_channels, args.model_z_channels,
encoder_core)
decoder_core = vgg.make_layers(vgg.decoder_cfg[args.model_type],
nonlinearity=nonlinearity,
nonlinearity_kwargs=kwargs)
decoder = knn.Unit(args.model_z_channels, args.model_in_channels,
decoder_core)
auto_encoder = autoencoder.AutoEncoder(encoder,
decoder,
init_weights=args.load is None).to(
args.device)
if args.load is not None:
auto_encoder.load(args.load)
""" optimizer """
optim = Adam(auto_encoder.parameters(), lr=args.lr)
""" variables """
best_loss = 100.0
run_dir = f'data/models/mapper/{args.model_type}/run_{args.run_id}'
writer = SummaryWriter(log_dir=run_dir)
global_step = 0
""" data """
dataset = ds.MapperDataset(32, 32, args.model_keypoints, args.dataset_train_len + args.dataset_test_len)
train, test = ds.random_split(dataset, [args.dataset_train_len, args.dataset_test_len])
train_l = DataLoader(train, batch_size=args.batch_size, shuffle=True, drop_last=True, pin_memory=True, num_workers=4)
test_l = DataLoader(test, batch_size=args.batch_size, shuffle=True, drop_last=True, pin_memory=True)
""" model """
encoder_core = vgg.make_layers(vgg.vgg_cfg[args.model_type])
encoder = knn.Unit(args.model_keypoints, 1, encoder_core)
mapper = transporter.TransporterMap(mapper=encoder, init_weights=args.load is None).to(args.device)
if args.load is not None:
mapper.load(args.load)
""" optimizer """
optim = Adam(mapper.parameters(), lr=1e-4)
""" apex mixed precision """
if args.device != 'cpu':
model, optimizer = amp.initialize(mapper, optim, opt_level=args.opt_level)
""" loss function """
criterion = nn.MSELoss()