pickle_path = os.path.join( './model', name, 'train_features' + str(i) + '.pkl') with open(pickle_path, "rb") as f: temp = pickle.load(f) features = torch.cat((features, temp)) features = features.view(-1, features.size(2)) print("Feature extraction finished") print( "Feeding extracted feature vectors to KMeans clustering model to generate clusters" ) # start clustering with features vector centers_path = train_cluster(features, centers_path) print("Clusters generated") model = model.convert_to_rpp_cluster() model.avgpool.centers_path = centers_path model = model.cuda() print(model) # step4: whole net training # print("STARTING STEP 4 OF PCB:") stage = 'full' full_train(model, criterion, stage, opt.epochs) else: model = PCB(len(class_names), num_bottleneck=256, num_parts=opt.parts, parts_ver=opt.PCB_Ver, checkerboard=opt.CB, share_conv=opt.share_conv)
elif opt.use_ftnet: model_structure = ft_net(opt.nclasses, stride=opt.stride) cluster = False if opt.PCB: model_structure = PCB(opt.nclasses, num_bottleneck=256, num_parts=opt.parts, parts_ver=opt.PCB_Ver, checkerboard=opt.CB, share_conv=opt.share_conv) if opt.RPP: model_structure = model_structure.convert_to_rpp() if opt.cluster: model_structure = model_structure.convert_to_rpp_cluster() print(model_structure) model_structure.avgpool.cluster_plots = opt.cluster_plots centers_path = os.path.join('./model', name, 'centers.pkl') model_structure.avgpool.centers_path = centers_path cluster = True #if opt.fp16: # model_structure = network_to_half(model_structure) model = load_network(model_structure) # Remove the final fc layer and classifier layer if opt.PCB: #if opt.fp16: # model = PCB_test(model[1])