num_workers=args.num_workers,
pin_memory=True,
collate_fn=my_collate)
test_loader = DataLoader(test_data, batch_size=32, collate_fn=my_collate)
print('Real: {0} \t Render: {1} \t Test: {2}'.format(len(real_loader),
len(render_loader),
len(test_loader)))
if np.isinf(args.max_iterations):
max_iterations = min(len(real_loader), len(render_loader))
else:
max_iterations = args.max_iterations
# my_model
if not args.multires:
model = OneBinDeltaModel(args.feature_network, num_classes, num_clusters,
args.N0, args.N1, args.N2, ndim)
else:
model = ProbabilisticOneDeltaPerBinModel(args.feature_network, num_classes,
num_clusters, args.N0, args.N1,
args.N2, args.N3, ndim)
# print(model)
# loss and optimizer
optimizer = optim.Adam(model.parameters(), lr=args.init_lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
# store stuff
writer = SummaryWriter(log_dir)
count = 0
val_loss = []
collate_fn=my_collate)
render_loader = DataLoader(render_data,
batch_size=args.num_workers,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
collate_fn=my_collate)
test_loader = DataLoader(test_data, batch_size=32)
print('Real: {0} \t Render: {1} \t Test: {2}'.format(len(real_loader),
len(render_loader),
len(test_loader)))
max_iterations = min(len(real_loader), len(render_loader))
# my_model
if not args.multires:
orig_model = OneBinDeltaModel(args.feature_network, num_classes,
num_clusters, N0, N1, N2, ndim)
else:
orig_model = OneDeltaPerBinModel(args.feature_network, num_classes,
num_clusters, N0, N1, N2, N3, ndim)
class JointCatPoseModel(nn.Module):
def __init__(self, oracle_model):
super().__init__()
# old stuff
self.num_classes = oracle_model.num_classes
self.num_clusters = oracle_model.num_clusters
self.ndim = oracle_model.ndim
self.feature_model = oracle_model.feature_model
self.bin_models = oracle_model.bin_models
self.res_models = oracle_model.res_models