log_dir = run_dir + '/' + args.name + '_stage_1'
create_folder(log_dir)
optimizer = optim.Adam(cnet.parameters(), lr=args.lr, weight_decay=args.weight_decay)
train_dataset = SingleImgDataset(args.train_path, scale_aug=False, rot_aug=False, num_models=n_models_train, num_views=args.num_views)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)
val_dataset = SingleImgDataset(args.val_path, scale_aug=False, rot_aug=False, test_mode=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=0)
print('num_train_files: '+str(len(train_dataset.filepaths)))
print('num_val_files: '+str(len(val_dataset.filepaths)))
trainer = ModelNetTrainer(cnet, train_loader, val_loader, optimizer, nn.CrossEntropyLoss(), 'svcnn', log_dir, num_views=1)
trainer.train(30)
elif args.svcnn_training_mode == 'load':
cnet.cuda()
log_dir = args.prefix + 'ckpt/svcnn/model-00025.pth'
# log_dir = ckpt_dir + '/20/mvcnn_stage_1/mvcnn/model-00025.pth'
model = torch.load(log_dir)
cnet.load_state_dict(model)
print('SVCNN trained model loaded!')
# STAGE 2
log_dir = run_dir + '/' + args.name + '_stage_2'
create_folder(log_dir)
cnet_2 = MVCNN(args.name, cnet, nclasses=40, cnn_name=args.cnn_name, num_views=args.num_views, constraint=args.constraint, w_m=args.w_m, T=args.T, preType=args.preType)
del cnet
new_params = list(cnet_2.main_net.parameters()) + list(cnet_2.main_net.parameters()) if args.preType != None else cnet_2.main_net.parameters()
params = new_params if args.freeze else cnet_2.parameters()
json.dump(vars(args), config_f)
config_f.close()
num_feature = 1024
# STAGE 1
log_dir = 'smtcloss_wonor_' + args.name + '_stage_1_' + localtime
create_folder(log_dir)
vgg = models.vgg11(pretrained=True)
cnet = SVCNN(args.name,
vgg,
num_feature,
nclasses=40,
pretraining=pretraining,
cnn_name=args.cnn_name)
if (torch.cuda.is_available()):
cnet = cnet.cuda()
print('use GPU to train ')
else:
print('don.t use gpu')
#有centor loss 时
center_loss = Triplet_Center_Loss()
softmax_loss = nn.CrossEntropyLoss()
optimizer_model = optim.SGD(cnet.parameters(),
lr=args.lr,
weight_decay=args.weight_decay,
momentum=0.9)
optimizer_centerloss = optim.SGD(center_loss.parameters(),
lr=args.lr_center)
# #triplet_loss
