def make_data(data, labels, output_dir, prefix): image_dir = prefix + '_images/' mkdir_if_missing(osp.join(output_dir, image_dir)) file_label_list = [] num = len(data) for i in xrange(num): img = np.rollaxis(data[i, :].reshape((3, 32, 32)), 0, 3) filename = '{:05d}.jpg'.format(i) imsave(osp.join(output_dir, image_dir, filename), img) file_label_list.append('{} {}'.format( osp.join(image_dir, filename), int(labels[i]))) write_list(file_label_list, osp.join(output_dir, prefix + '.txt'))
def main(args): mkdir_if_missing(args.output_dir) # training data data = [] labels = [] for i in xrange(1, 6): dic = unpickle(osp.join(args.data_root, 'data_batch_{}'.format(i))) data.append(dic['data']) labels = np.r_[labels, dic['labels']] data = np.vstack(data) make_data(data, labels, args.output_dir, 'train') # test data dic = unpickle(osp.join(args.data_root, 'test_batch')) make_data(dic['data'], dic['labels'], args.output_dir, 'test')
def _download_data(self): if osp.exists(self.root): print("This dataset has been downloaded.") return mkdir_if_missing(self.root) fpath = osp.join(self.root, osp.basename(self.dataset_url)) print("Downloading iLIDS-VID dataset") url_opener = urllib.URLopener() url_opener.retrieve(self.dataset_url, fpath) print("Extracting files") tar = tarfile.open(fpath) tar.extractall(path=self.root) tar.close()
def _download_data(self): if osp.exists(self.dataset_dir): print("This dataset has been downloaded.") return print("Creating directory {}".format(self.dataset_dir)) mkdir_if_missing(self.dataset_dir) fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url)) print("Downloading DukeMTMC-VideoReID dataset") url_opener = urllib.URLopener() url_opener.retrieve(self.dataset_url, fpath) print("Extracting files") zip_ref = zipfile.ZipFile(fpath, 'r') zip_ref.extractall(self.dataset_dir) zip_ref.close()
def main(args): output_dir = args.output_dir mkdir_if_missing(osp.join(args.output_dir, 'cam_0')) mkdir_if_missing(osp.join(args.output_dir, 'cam_1')) # Collect the person_id and view_id into dict images = glob(osp.join(args.input_dir, 'RGB', '*.bmp')) pdict = defaultdict(lambda: defaultdict(list)) for imname in images: pid, vid = osp.basename(imname).split('_')[0:2] pdict[pid][vid].append(imname) # Randomly choose half of the views as cam_0, others as cam_1 identities = [] for i, pid in enumerate(pdict): vids = pdict[pid].keys() num_views = len(vids) np.random.shuffle(vids) p_images = [[], []] for vid in vids[:(num_views // 2)]: for src_file in pdict[pid][vid]: tgt_file = 'cam_0/{:05d}_{:05d}.bmp'.format( i, len(p_images[0])) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[0].append(tgt_file) for vid in vids[(num_views // 2):]: for src_file in pdict[pid][vid]: tgt_file = 'cam_1/{:05d}_{:05d}.bmp'.format( i, len(p_images[1])) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[1].append(tgt_file) identities.append(p_images) # Save meta information into a json file meta = {'name': '3DPeS', 'shot': 'multiple', 'num_cameras': 2} meta['identities'] = identities write_json(meta, osp.join(args.output_dir, 'meta.json')) # Randomly create a training and test split num = len(identities) pids = np.random.permutation(num) trainval_pids = sorted(pids[100:]) test_pids = sorted(pids[:100]) split = { 'trainval': trainval_pids, 'test_probe': test_pids, 'test_gallery': test_pids } write_json(split, osp.join(output_dir, 'split.json'))
def main(args): meta = read_json(osp.join(args.dataset_dir, 'meta.json')) split = read_json(osp.join(args.dataset_dir, 'split.json')) identities = np.asarray(meta['identities']) # Make train / val. Consider single or multiple shot. trainval = identities[split['trainval']] if meta['shot'] == 'single': # When single shot, to ensure each identity has at least one training # image, we first randomly choose validation identities, then randomly # split their views equally for training and validation. num_val = int(len(trainval) * args.val_ratio) * 2 np.random.shuffle(trainval) train = list(trainval[num_val:]) val = [] for views in trainval[:num_val]: np.random.shuffle(views) train.append(views[:len(views) // 2]) val.append(views[len(views) // 2:]) train = _get_list(train) val = _get_list(val) else: # When multiple shots, we just randomly split the trainval images trainval = _get_list(trainval) np.random.shuffle(trainval) num_val = int(len(trainval) * args.val_ratio) train = trainval[num_val:] val = trainval[:num_val] # Make test probe / gallery. Probe identities should be a subset of # gallery's. First half views are probe, others are gallery. assert len(set(split['test_probe']) - set(split['test_gallery'])) == 0 test_probe, test_gallery = [], [] for views in identities[split['test_probe']]: test_probe.append(views[:len(views) // 2]) test_gallery.append(views[len(views) // 2:]) only_in_gallery = list( set(split['test_gallery']) - set(split['test_probe'])) test_gallery.extend(identities[only_in_gallery]) test_probe = _get_list(test_probe) test_gallery = _get_list(test_gallery) # Save to files mkdir_if_missing(args.output_dir) _save(train, osp.join(args.output_dir, 'train.txt')) _save(val, osp.join(args.output_dir, 'val.txt')) _save(test_probe, osp.join(args.output_dir, 'test_probe.txt')) _save(test_gallery, osp.join(args.output_dir, 'test_gallery.txt'))
def __init__(self, params): self.params = params self.starting_time = time.time() self.elapsed_time = 0 utils.mkdir_if_missing(self.params.MODEL_SAVE_DIR) self.val_datagen = datagen.SequenceDataGenerator( num_shape = self.params.NUM_SHAPE, image_size = self.params.IMG_SIZE, sequence_len = self.params.SEQUENCE_LEN, random_size = True, rotate_shapes = self.params.ROTATE_SHAPES ) self.train_data_loader = dataset.SequenceDataLoader( dataset_path=self.params.TRAIN_SET_PATH, shuffle=True) self.val_data_loader = dataset.SequenceDataLoader( dataset_path=self.params.VAL_SET_PATH, shuffle=False) self.test_data_loader = dataset.SequenceDataLoader( dataset_path=self.params.TEST_SET_PATH, shuffle=False)
def train_teacher(): """ This function trains a teacher (teacher id) among an ensemble of nb_teachers models for the dataset specified. :param dataset: string corresponding to dataset (svhn, cifar10) :param nb_teachers: total number of teachers in the ensemble :param teacher_id: id of the teacher being trained :return: True if everything went well """ # If working directories do not exist, create them #assert utils.mkdir_if_missing(config.data_dir) #assert utils.mkdir_if_misshing(config.train_dir) print("Initializing dataset {}".format(config.dataset)) dataset = data_manager.init_img_dataset( root=config.data_dir, name=config.dataset, ) # Load the dataset for i in range(0,config.nb_teachers): # Retrieve subset of data for this teacher if config.dataset == 'celeba': data, labels = dataset._data_partition(config.nb_teachers,i) print("Length of training data: " + str(len(data))) # Define teacher checkpoint filename and full path print('data.shape for each teacher') dir_path = os.path.join(config.save_model,'pate_'+config.dataset+str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) #filename = os.path.join(dir_path, str(config.nb_teachers) + '_teachers_' + str(i) + '_resnet.checkpoint.pth.tar') filename = os.path.join(dir_path, str(config.nb_teachers) + '_teachers_' + str(i) + config.arch+'.checkpoint.pth.tar') print('save_path for teacher{} is {}'.format(i,filename)) network.train_each_teacher(config.teacher_epoch,data, labels, dataset.test_data, dataset.test_label, filename) return True
def main(args): mkdir_if_missing(args.output_dir) # training data data = [] labels = [] for i in xrange(1, 6): dic = unpickle(osp.join(args.data_root, 'data_batch_{}'.format(i))) data.append(dic['data']) labels = np.r_[labels, dic['labels']] data = np.vstack(data) make_data(data, labels, args.output_dir, 'train') # test data dic = unpickle(osp.join(args.data_root, 'test_batch')) make_data(dic['data'], dic['labels'], args.output_dir, 'test') # Identity for confusion initialization matrix_I = np.identity(10) write_matrix(matrix_I, osp.join(args.output_dir, 'identity.txt')) pickle(matrix_I, osp.join(args.output_dir, 'identity.pkl'))
def train_tracher(): """ Partition the entire private (training) data into config.nb_teacher subsets and train config.nb_teacher teacher models """ # Load the dataset if config.dataset == 'mnist': train_dataset = dataset.MNIST(root=config.data_dir, train=True, download=True) test_dataset = dataset.MNIST(root=config.data_dir, train=False, download=True) ori_train_data= [ data[0] for idx, data in enumerate(train_dataset)] ori_test_data = [ data[0] for idx, data in enumerate(test_dataset)] test_labels = test_dataset.targets train_labels = train_dataset.targets elif config.dataset =='svhn': train_dataset = dataset.SVHN(root=config.data_dir, split='train', download=True) extra_dataset = dataset.SVHN(root=config.data_dir, split='extra', download=True) test_dataset = dataset.SVHN(root=config.data_dir, split='test', download=True) ori_train_data = np.concatenate((train_dataset.data,extra_dataset.data),axis=0) print('ori data shape', ori_train_data.shape) ori_train_data = np.transpose(ori_train_data, (0, 2, 3, 1)) print('orig data shape', ori_train_data.shape) #ori_train_data= [ data[0] for idx, data in enumerate(train_dataset.data)] #for data in extra_dataset.data: # ori_train_data.append(data) #ori_test_data = [ data[0] for idx, data in enumerate(test_dataset.data)] ori_test_data = np.transpose(test_dataset.data, (0,2,3,1)) test_labels = test_dataset.labels extra_labels = extra_dataset.labels train_labels = [ll for ll in train_dataset.labels] for ll in extra_labels: train_labels.append(ll) batch_len = int(len(ori_train_data)/config.nb_teachers) for i in range(0,1): dir_path = os.path.join(config.save_model,'pate_'+str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) filename = os.path.join(dir_path, str(config.nb_teachers) + '_teachers_' + str(i) + config.arch+'.checkpoint.pth.tar') print('save_path for teacher{} is {}'.format(i,filename)) start = i * batch_len end = (i + 1) * batch_len t_data = ori_train_data[start : end] t_labels = train_labels[start: end] network.train_each_teacher(config.teacher_epoch, t_data, t_labels, ori_test_data, test_labels, filename)
def main(args): id_offset = 0 merged_train_kv = {} merged_val_kv = {} for dataset_dir, db_dir in zip(args.dataset_dirs, args.db_dirs): train_files, train_labels = read_kv(osp.join(db_dir, 'train.txt')) val_files, val_labels = read_kv(osp.join(db_dir, 'val.txt')) unique_ids = set(map(int, train_labels + val_labels)) id_mapping = {idx: i + id_offset for i, idx in enumerate(unique_ids)} for k, v in zip(train_files, train_labels): merged_train_kv[osp.join(dataset_dir, k)] = id_mapping[int(v)] for k, v in zip(val_files, val_labels): merged_val_kv[osp.join(dataset_dir, k)] = id_mapping[int(v)] id_offset += len(id_mapping) mkdir_if_missing(osp.join(args.output_dir)) train_list = [k + ' ' + str(v) for k, v in merged_train_kv.iteritems()] np.random.shuffle(train_list) write_list(train_list, osp.join(args.output_dir, 'train.txt')) write_kv(merged_val_kv.keys(), map(str, merged_val_kv.values()), osp.join(args.output_dir, 'val.txt')) print "Max ID:", id_offset
def main(args): output_dir = args.output_dir mkdir_if_missing(osp.join(args.output_dir, 'cam_0')) mkdir_if_missing(osp.join(args.output_dir, 'cam_1')) # Collect the images of each person into dict images = glob(osp.join(args.ilids_dir, 'Persons', '*.jpg')) pdict = defaultdict(list) for imname in images: pid = int(osp.basename(imname)[:4]) pdict[pid].append(imname) # Randomly choose half of the images as cam_0, others as cam_1 identities = [] for i, (pid, images) in enumerate(pdict.iteritems()): num = len(images) np.random.shuffle(images) p_images = [[], []] for src_file in images[:(num // 2)]: tgt_file = 'cam_0/{:05d}_{:05d}.jpg'.format(i, len(p_images[0])) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[0].append(tgt_file) for src_file in images[(num // 2):]: tgt_file = 'cam_1/{:05d}_{:05d}.jpg'.format(i, len(p_images[1])) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[1].append(tgt_file) identities.append(p_images) # Save meta information into a json file meta = {'name': 'i-LIDS', 'shot': 'multiple', 'num_cameras': 2} meta['identities'] = identities write_json(meta, osp.join(output_dir, 'meta.json')) # Randomly create a training and test split num = len(identities) pids = np.random.permutation(num) trainval_pids = sorted(pids[:num // 2]) test_pids = sorted(pids[num // 2:]) split = { 'trainval': trainval_pids, 'test_probe': test_pids, 'test_gallery': test_pids } write_json(split, osp.join(output_dir, 'split.json'))
def ensemble_preds(nb_teachers, stdnt_data): """ Given a dataset, a number of teachers, and some input data, this helper function queries each teacher for predictions on the data and returns all predictions in a single array. (That can then be aggregated into one single prediction per input using aggregation.py (cf. function prepare_student_data() below) :param dataset: string corresponding to mnist, cifar10, or svhn :param nb_teachers: number of teachers (in the ensemble) to learn from :param stdnt_data: unlabeled student training data :return: 3d array (teacher id, sample id, probability per class) """ result_shape = (nb_teachers, len(stdnt_data), config.nb_labels) # Create array that will hold result result = np.zeros(result_shape, dtype=np.float32) # Get predictions from each teacher for teacher_id in range(nb_teachers): # Compute path of checkpoint file for teacher model with ID teacher_id if config.dataset == 'celeba': dir_path = os.path.join( config.save_model, 'pate_num_teacher_' + str(config.nb_teachers)) elif config.dataset == 'market': dir_path = os.path.join( config.save_model, 'pate_' + config.dataset + str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) filename = os.path.join( dir_path, str(config.nb_teachers) + '_teachers_' + str(teacher_id) + config.arch + '.checkpoint.pth.tar') result[teacher_id] = network.pred(stdnt_data, filename) # This can take a while when there are a lot of teachers so output status print("Computed Teacher " + str(teacher_id) + " softmax predictions") return result
def gen_dataset(self, params, dataset_type, seed=0): random.seed(seed) np.random.seed(seed) if dataset_type == 'train': num_seq = params.TRAIN_NUM_SEQ path = params.TRAIN_SET_PATH elif dataset_type == 'val': num_seq = params.VAL_NUM_SEQ path = params.VAL_SET_PATH elif dataset_type == 'test': num_seq = params.TEST_NUM_SEQ path = params.TEST_SET_PATH else: raise ValueError('dataset_type must be train, val, or test') for i in range(num_seq): sdg = SequenceDataGenerator(params.NUM_SHAPE, params.IMG_SIZE, params.SEQUENCE_LEN, params.RANDOM_SIZE, params.ROTATE_SHAPES) utils.update_progress(i / num_seq) seq = sdg.get_sequence() pickle_folder_path = os.path.join(path, f'seq_{i}') utils.mkdir_if_missing(pickle_folder_path) pickle_full_path = os.path.join(pickle_folder_path, 'sequence.pickle') with open(pickle_full_path, 'wb') as handle: pickle.dump(seq, handle) image_count = 0 for info in seq: image = info['image'] image_folder_path = os.path.join(path, f'seq_{i}', 'images') utils.mkdir_if_missing(image_folder_path) image_full_path = os.path.join(image_folder_path, f'{image_count:05d}.png') image = (image * 255).astype(np.uint8) cv2.imwrite(image_full_path, image) image_count += 1
def main(args): # cam_0 to cam_15 for i in xrange(16): mkdir_if_missing(osp.join(args.output_dir, 'cam_' + str(i))) images = glob(osp.join(args.shinpuhkan_dir, 'images', '*.jpg')) images.sort() identities = [] prev_pid = -1 for name in images: name = osp.basename(name) p_id = int(name[0:3]) - 1 c_id = int(name[4:6]) - 1 if prev_pid != p_id: identities.append([]) prev_cid = -1 p_images = identities[-1] if prev_cid != c_id: p_images.append([]) v_images = p_images[-1] file_name = 'cam_{}/{:05d}_{:05d}.jpg'.format(c_id, p_id, len(v_images)) shutil.copy(osp.join(args.shinpuhkan_dir, 'images', name), osp.join(args.output_dir, file_name)) v_images.append(file_name) prev_pid = p_id prev_cid = c_id # Save meta information into a json file meta = {'name': 'Shinpuhkan', 'shot': 'multiple', 'num_cameras': 16} meta['identities'] = identities write_json(meta, osp.join(args.output_dir, 'meta.json')) # We don't test on this dataset. Just use all the data for train / val. split = { 'trainval': range(len(identities)), 'test_probe': [], 'test_gallery': [] } write_json(split, osp.join(args.output_dir, 'split.json'))
def main(args): output_dir = args.output_dir mkdir_if_missing(osp.join(args.output_dir, 'cam_0')) mkdir_if_missing(osp.join(args.output_dir, 'cam_1')) identities = [] cam1_images = glob(osp.join(args.viper_dir, 'cam_a', '*.bmp')) cam2_images = glob(osp.join(args.viper_dir, 'cam_b', '*.bmp')) cam1_images.sort() cam2_images.sort() assert len(cam1_images) == len(cam2_images) for i in xrange(len(cam1_images)): p_id = len(identities) p_images = [] # view-0 file_name = 'cam_0/{:05d}_{:05d}.bmp'.format(p_id, 0) shutil.copy(cam1_images[i], osp.join(args.output_dir, file_name)) p_images.append([file_name]) # view-1 file_name = 'cam_1/{:05d}_{:05d}.bmp'.format(p_id, 0) shutil.copy(cam2_images[i], osp.join(args.output_dir, file_name)) p_images.append([file_name]) identities.append(p_images) # Save meta information into a json file meta = {'name': 'VIPeR', 'shot': 'single', 'num_cameras': 2} meta['identities'] = identities write_json(meta, osp.join(args.output_dir, 'meta.json')) # Randomly create a training and test split num = len(identities) pids = np.random.permutation(num) trainval_pids = sorted(pids[:num // 2]) test_pids = sorted(pids[num // 2:]) split = {'trainval': trainval_pids, 'test_probe': test_pids, 'test_gallery': test_pids} write_json(split, osp.join(output_dir, 'split.json'))
def train_student( nb_teachers): """ This function trains a student using predictions made by an ensemble of teachers. The student and teacher models are trained using the same neural network architecture. :param dataset: string corresponding to celeba :param nb_teachers: number of teachers (in the ensemble) to learn from :return: True if student training went well """ # Call helper function to prepare student data using teacher predictions stdnt_dataset = prepare_student_data(nb_teachers, save=True) # Unpack the student dataset stdnt_data, stdnt_labels, stdnt_test_data, stdnt_test_labels = stdnt_dataset dir_path = os.path.join(config.save_model, config.dataset) dir_path = os.path.join(dir_path, 'knn_num_neighbor_' + str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) if config.resnet: filename = os.path.join(dir_path,str(config.nb_teachers) + '_stdnt_resnet.checkpoint.pth.tar') print('stdnt_label used for train', stdnt_labels.shape) network.train_each_teacher(config.student_epoch, stdnt_data, stdnt_labels, stdnt_test_data, stdnt_test_labels, filename) return True
def main(args): output_dir = args.output_dir mkdir_if_missing(osp.join(args.output_dir, 'cam_0')) mkdir_if_missing(osp.join(args.output_dir, 'cam_1')) # Randomly choose 100 people from the 200 shared people as test probe p = list(np.random.permutation(200)) test_probe = range(100) test_gallery = range(100) identities = [] for pid in p[:100]: p_images = [] src_file = osp.join(args.prid_dir, 'single_shot', 'cam_a', 'person_{:04d}.png'.format(pid + 1)) tgt_file = osp.join('cam_0', '{:05d}_00000.png'.format(len(identities))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images.append([tgt_file]) src_file = osp.join(args.prid_dir, 'single_shot', 'cam_b', 'person_{:04d}.png'.format(pid + 1)) tgt_file = osp.join('cam_1', '{:05d}_00000.png'.format(len(identities))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images.append([tgt_file]) identities.append(p_images) # Other 100 people from the 200 as a part of trainval # Choose 10 images randomly from the multi-shot images trainval = range(100, 200) for pid in p[100:]: p_images = [[], []] images = glob( osp.join(args.prid_dir, 'multi_shot', 'cam_a', 'person_{:04d}'.format(pid + 1), '*.png')) images = np.random.choice(images, size=min(10, len(images)), replace=False) for src_file in images: tgt_file = osp.join( 'cam_0', '{:05d}_{:05d}.png'.format(len(identities), len(p_images[0]))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[0].append(tgt_file) images = glob( osp.join(args.prid_dir, 'multi_shot', 'cam_b', 'person_{:04d}'.format(pid + 1), '*.png')) images = np.random.choice(images, size=min(10, len(images)), replace=False) for src_file in images: tgt_file = osp.join( 'cam_1', '{:05d}_{:05d}.png'.format(len(identities), len(p_images[1]))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[1].append(tgt_file) identities.append(p_images) # 201 to 385 cam_a people as another part of trainval for pid in xrange(200, 385): p_images = [[], []] images = glob( osp.join(args.prid_dir, 'multi_shot', 'cam_a', 'person_{:04d}'.format(pid + 1), '*.png')) images = np.random.choice(images, size=min(10, len(images)), replace=False) for src_file in images: tgt_file = osp.join( 'cam_0', '{:05d}_{:05d}.png'.format(len(identities), len(p_images[0]))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images[0].append(tgt_file) trainval.append(len(identities)) identities.append(p_images) # 201 to 749 cam_b people as additional test gallery for pid in xrange(200, 749): src_file = osp.join(args.prid_dir, 'single_shot', 'cam_b', 'person_{:04d}.png'.format(pid + 1)) tgt_file = osp.join('cam_1', '{:05d}_00000.png'.format(len(identities))) shutil.copy(src_file, osp.join(args.output_dir, tgt_file)) p_images = [[], [tgt_file]] test_gallery.append(len(identities)) identities.append(p_images) # Save meta information into a json file meta = {'name': 'PRID', 'shot': 'multiple', 'num_cameras': 2} meta['identities'] = identities write_json(meta, osp.join(args.output_dir, 'meta.json')) # We have only one split split = { 'trainval': trainval, 'test_probe': test_probe, 'test_gallery': test_gallery } write_json(split, osp.join(output_dir, 'split.json'))
def train_task(args, train_loader, current_task, prototype={}, pre_index=0): num_class_per_task = (args.num_class - args.nb_cl_fg) // args.num_task task_range = list( range(args.nb_cl_fg + (current_task - 1) * num_class_per_task, args.nb_cl_fg + current_task * num_class_per_task)) if num_class_per_task == 0: pass # JT else: old_task_factor = args.nb_cl_fg // num_class_per_task + current_task - 1 print(old_task_factor) log_dir = os.path.join(args.ckpt_dir, args.log_dir) mkdir_if_missing(log_dir) sys.stdout = logging.Logger( os.path.join(log_dir, 'log_task{}.txt'.format(current_task))) tb_writer = SummaryWriter(log_dir) display(args) if 'miniimagenet' in args.data: model = models.create('resnet18_imagenet', pretrained=False, feat_dim=args.feat_dim, embed_dim=args.num_class, hidden_dim=256, norm=True) elif 'cifar100' in args.data: model = models.create('resnet18_cifar', pretrained=False, feat_dim=args.feat_dim, hidden_dim=256, embed_dim=args.num_class, norm=True) if current_task > 0: if 'miniimagenet' in args.data: model = models.create('resnet18_imagenet', pretrained=False, feat_dim=args.feat_dim, embed_dim=args.num_class, hidden_dim=256, norm=True) elif 'cifar100' in args.data: model = models.create('resnet18_cifar', pretrained=False, feat_dim=args.feat_dim, hidden_dim=256, embed_dim=args.num_class, norm=True) model = torch.load( os.path.join( log_dir, 'task_' + str(current_task - 1).zfill(2) + '_%d_model.pkl' % int(args.epochs - 1))) model_old = deepcopy(model) model_old.eval() model_old = freeze_model(model_old) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # model = model.cuda() model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = StepLR(optimizer, step_size=args.lr_decay_step, gamma=args.lr_decay) loss_mse = torch.nn.MSELoss(reduction='sum') # # Loss weight for gradient penalty used in W-GAN lambda_gp = args.lambda_gp lambda_lwf = args.gan_tradeoff # Initialize generator and discriminator if current_task == 0: generator = Generator(feat_dim=args.feat_dim, latent_dim=args.latent_dim, hidden_dim=args.hidden_dim, class_dim=args.num_class, norm=True) discriminator = Discriminator(feat_dim=args.feat_dim, hidden_dim=args.hidden_dim, class_dim=args.num_class) else: generator = torch.load( os.path.join( log_dir, 'task_' + str(current_task - 1).zfill(2) + '_%d_model_generator.pkl' % int(args.epochs_gan - 1))) discriminator = torch.load( os.path.join( log_dir, 'task_' + str(current_task - 1).zfill(2) + '_%d_model_discriminator.pkl' % int(args.epochs_gan - 1))) generator_old = deepcopy(generator) generator_old.eval() generator_old = freeze_model(generator_old) cuda = torch.cuda.is_available() FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor g_len = 0 d_len = 0 for p in generator.parameters(): g_len += 1 for p in discriminator.parameters(): d_len += 1 learned_lrs = [] params = [] for i in range(args.update_step): g_lrs = [ Variable(FloatTensor(1).fill_(args.update_lr), requires_grad=True) ] * g_len # len(generator.parameters()) d_lrs = [ Variable(FloatTensor(1).fill_(args.update_lr), requires_grad=True) ] * d_len # len(discriminator.parameters()) learned_lrs.append((g_lrs, d_lrs)) for param_list in learned_lrs[i]: params += param_list generator = generator.to(device) discriminator = discriminator.to(device) optimizer_G = torch.optim.Adam(generator.parameters(), lr=args.gan_lr, betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=args.gan_lr, betas=(0.5, 0.999)) optimizer_lr = torch.optim.Adam(params, lr=args.lr) scheduler_G = StepLR(optimizer_G, step_size=150, gamma=0.3) scheduler_D = StepLR(optimizer_D, step_size=150, gamma=0.3) for p in generator.parameters(): # set requires_grad to False p.requires_grad = False for epoch in range(args.epochs): loss_log = { 'C/loss': 0.0, 'C/loss_aug': 0.0, 'C/loss_cls': 0.0, 'C/loss_cls_q': 0.0 } ##### MAML on feature extraction for step, (x_spt, y_spt, x_qry, y_qry) in enumerate(train_loader): x_spt, y_spt, x_qry, y_qry = x_spt.to(device), y_spt.to( device), x_qry.to(device), y_qry.to(device) loss = torch.zeros(1).to(device) loss_cls = torch.zeros(1).to(device) loss_aug = torch.zeros(1).to(device) loss_tmp = torch.zeros(1).to(device) BatchSize, setsz, c_, h, w = x_spt.size() querysz = x_qry.size(1) losses_q = [0 for _ in range(args.update_step + 1) ] # losses_q[i] is the loss on step i corrects = [0.0 for _ in range(args.update_step + 1)] correct_s = [0.0 for _ in range(args.update_step + 1)] y_onehot = torch.cuda.FloatTensor(setsz, args.num_class) y_onehot_q = torch.cuda.FloatTensor(querysz, args.num_class) for i in range(args.BatchSize): # 1. run the i-th task and compute loss for k=0 embed_feat = model(x_spt[i]) if current_task == 0: soft_feat = model.embed(embed_feat) loss_cls = torch.nn.CrossEntropyLoss()(soft_feat, y_spt[i]) loss = loss.clone() + loss_cls else: embed_feat_old = model_old(x_spt[i]) ### Feature Extractor Loss if current_task > 0: loss_aug = torch.dist(embed_feat, embed_feat_old, 2) loss = loss.clone( ) + args.tradeoff * loss_aug * old_task_factor ### Replay and Classification Loss if current_task > 0: embed_sythesis = [] embed_label_sythesis = [] ind = list(range(len(pre_index))) if args.mean_replay: for _ in range(setsz): np.random.shuffle(ind) tmp = prototype['class_mean'][ ind[0]] + np.random.normal( ) * prototype['class_std'][ind[0]] embed_sythesis.append(tmp) embed_label_sythesis.append( prototype['class_label'][ind[0]]) embed_sythesis = np.asarray(embed_sythesis) embed_label_sythesis = np.asarray(embed_label_sythesis) embed_sythesis = torch.from_numpy(embed_sythesis).to( device) embed_label_sythesis = torch.from_numpy( embed_label_sythesis) else: for _ in range(setsz): np.random.shuffle(ind) embed_label_sythesis.append(pre_index[ind[0]]) embed_label_sythesis = np.asarray(embed_label_sythesis) embed_label_sythesis = torch.from_numpy( embed_label_sythesis).to(device) y_onehot.zero_() y_onehot.scatter(1, embed_label_sythesis[:, None], 1) syn_label_pre = y_onehot.to(device) z = torch.Tensor( np.random.normal( 0, 1, (setsz, args.latent_dim))).to(device) embed_sythesis = generator(z, syn_label_pre) embed_sythesis = torch.cat((embed_feat, embed_sythesis)) embed_label_sythesis = torch.cat( (y_spt[i], embed_label_sythesis.to(device))) soft_feat_syt = model.embed(embed_sythesis) batch_size1 = x_spt[i].shape[0] batch_size2 = embed_feat.shape[0] loss_cls = torch.nn.CrossEntropyLoss()( soft_feat_syt[:batch_size1], embed_label_sythesis[:batch_size1]) loss_cls_old = torch.nn.CrossEntropyLoss()( soft_feat_syt[batch_size2:], embed_label_sythesis[batch_size2:]) loss_cls += loss_cls_old * old_task_factor loss_cls /= args.nb_cl_fg // num_class_per_task + current_task loss += loss_cls grad = torch.autograd.grad(loss, model.parameters(), create_graph=True, retain_graph=True) # this is the loss and accuracy before first update with torch.no_grad(): # [setsz, nway] embed_feat_q = model(x_qry[i]) soft_feat_q = model.embed(embed_feat_q) loss_q = torch.nn.CrossEntropyLoss()(soft_feat_q, y_qry[i]) losses_q[0] += loss_q embed_feat = model(x_spt[i]) soft_feat = model.embed(embed_feat) pred_s = F.softmax(soft_feat, dim=1).argmax(dim=1) corr = torch.eq(pred_s, y_spt[i]).sum().item() # convert to numpy correct_s[0] = correct_s[0] + corr pred_q = F.softmax(soft_feat_q, dim=1).argmax(dim=1) correct = torch.eq(pred_q, y_qry[i]).sum().item() corrects[0] = corrects[0] + correct # this is the loss and accuracy after the first update with torch.no_grad(): # [setsz, nway] for e, param in enumerate(model.parameters(), 0): param.data -= args.update_lr * grad[e] embed_feat_q = model(x_qry[i]) soft_feat_q = model.embed(embed_feat_q) loss_q = torch.nn.CrossEntropyLoss()(soft_feat_q, y_qry[i]) losses_q[1] += loss_q # [setsz] embed_feat = model(x_spt[i]) soft_feat = model.embed(embed_feat) pred_s = F.softmax(soft_feat, dim=1).argmax(dim=1) corr = torch.eq(pred_s, y_spt[i]).sum().item() # convert to numpy correct_s[1] = correct_s[1] + corr pred_q = F.softmax(soft_feat_q, dim=1).argmax(dim=1) correct = torch.eq(pred_q, y_qry[i]).sum().item() corrects[1] = corrects[1] + correct for k in range(1, args.update_step): # 1. run the i-th task and compute loss for k=1~K-1 embed_feat = model(x_spt[i]) loss = torch.zeros(1).to(device) if current_task > 0: embed_feat_old = model_old(x_spt[i]) loss_aug = torch.dist(embed_feat, embed_feat_old, 2) loss += args.tradeoff * loss_aug * old_task_factor embed_sythesis = [] embed_label_sythesis = [] ind = list(range(len(pre_index))) if args.mean_replay: for _ in range(setsz): np.random.shuffle(ind) tmp = prototype['class_mean'][ ind[0]] + np.random.normal( ) * prototype['class_std'][ind[0]] embed_sythesis.append(tmp) embed_label_sythesis.append( prototype['class_label'][ind[0]]) embed_sythesis = np.asarray(embed_sythesis) embed_label_sythesis = np.asarray( embed_label_sythesis) embed_sythesis = torch.from_numpy( embed_sythesis).to(device) embed_label_sythesis = torch.from_numpy( embed_label_sythesis) else: for _ in range(setsz): np.random.shuffle(ind) embed_label_sythesis.append(pre_index[ind[0]]) embed_label_sythesis = np.asarray( embed_label_sythesis) embed_label_sythesis = torch.from_numpy( embed_label_sythesis).to(device) y_onehot.zero_() y_onehot.scatter(1, embed_label_sythesis[:, None], 1) syn_label_pre = y_onehot.to(device) z = torch.Tensor( np.random.normal( 0, 1, (setsz, args.latent_dim))).to(device) embed_sythesis = generator(z, syn_label_pre) embed_sythesis = torch.cat( (embed_feat, embed_sythesis)) embed_label_sythesis = torch.cat( (y_spt[i], embed_label_sythesis.to(device))) soft_feat_syt = model.embed(embed_sythesis) batch_size1 = x_spt[i].shape[0] batch_size2 = embed_feat.shape[0] loss_cls = torch.nn.CrossEntropyLoss()( soft_feat_syt[:batch_size1], embed_label_sythesis[:batch_size1]) loss_cls_old = torch.nn.CrossEntropyLoss()( soft_feat_syt[batch_size2:], embed_label_sythesis[batch_size2:]) loss_cls += loss_cls_old * old_task_factor loss_cls /= args.nb_cl_fg // num_class_per_task + current_task loss += loss_cls else: soft_feat = model.embed(embed_feat) loss_cls = torch.nn.CrossEntropyLoss()(soft_feat, y_spt[i]) loss += loss_cls # 2. compute grad on theta_pi grad = torch.autograd.grad(loss, model.parameters(), create_graph=True, retain_graph=True, allow_unused=True) # 3. theta_pi = theta_pi - train_lr * grad for e, param in enumerate(model.parameters(), 0): param.data -= args.update_lr * grad[e] embed_feat = model(x_spt[i]) soft_feat = model.embed(embed_feat) # soft_feat = mlp(embed_feat) embed_feat_q = model(x_qry[i]) soft_feat_q = model.embed(embed_feat_q) # loss_q will be overwritten and just keep the loss_q on last update step. loss_q = torch.nn.CrossEntropyLoss()(soft_feat_q, y_qry[i]) losses_q[k + 1] += loss_q with torch.no_grad(): pred_s = F.softmax(soft_feat, dim=1).argmax(dim=1) corr = torch.eq( pred_s, y_spt[i]).sum().item() # convert to numpy correct_s[k + 1] = correct_s[k + 1] + corr pred_q = F.softmax(soft_feat_q, dim=1).argmax(dim=1) correct = torch.eq( pred_q, y_qry[i]).sum().item() # convert to numpy corrects[k + 1] = corrects[k + 1] + correct # end of all tasks # sum over all losses on query set across all tasks loss_q = losses_q[-1] / BatchSize loss_q = Variable(loss_q, requires_grad=True) # optimize theta parameters optimizer.zero_grad() loss_q.backward() optimizer.step() scheduler.step() accs = np.array([float(c) for c in corrects]) / float(querysz * BatchSize) accs_spt = np.array([float(c) for c in correct_s]) / float( setsz * BatchSize) loss_log['C/loss'] += loss.item() loss_log['C/loss_cls'] += loss_cls.item() loss_log['C/loss_aug'] += args.tradeoff * loss_aug.item( ) if args.tradeoff != 0 else 0 loss_log['C/loss_cls_q'] += loss_q.item() del loss_cls del loss_q print( '[Metric Epoch %05d]\t Total Loss: %.3f \t LwF Loss: %.3f \t Spt Accuracy FeatureX: %.3f \t Query Loss: %.3f \t Query Accuracy FeatureX: %.3f \t' % (epoch + 1, loss_log['C/loss'], loss_log['C/loss_aug'], accs_spt[-1], loss_log['C/loss_cls_q'], accs[-1])) for k, v in loss_log.items(): if v != 0: tb_writer.add_scalar( 'Task {} - Classifier/{}'.format(current_task, k), v, epoch + 1) tb_writer.add_scalar('Task {}'.format(current_task), accs[-1], epoch + 1) if epoch == args.epochs - 1: torch.save( model, os.path.join( log_dir, 'task_' + str(current_task).zfill(2) + '_%d_model.pkl' % epoch)) ################# feature extraction training end ######################## ############################################## GAN Training #################################################### model = model.eval() for p in model.parameters(): # set requires_grad to False p.requires_grad = False for p in generator.parameters(): # set requires_grad to True p.requires_grad = True for p in discriminator.parameters(): p.requires_grad = True criterion_softmax = torch.nn.CrossEntropyLoss().to(device) if current_task != args.num_task: for epoch in range(args.epochs_gan): loss_log = { 'D/loss': 0.0, 'D/new_rf': 0.0, 'D/new_lbls': 0.0, 'D/new_gp': 0.0, 'D/prev_rf': 0.0, 'D/prev_lbls': 0.0, 'D/prev_gp': 0.0, 'D/loss_q': 0.0, 'D/new_rf_q': 0.0, 'D/new_lbls_q': 0.0, 'D/new_gp_q': 0.0, 'G/loss': 0.0, 'G/new_rf': 0.0, 'G/new_lbls': 0.0, 'G/prev_rf': 0.0, 'G/prev_mse': 0.0, 'G/new_classifier': 0.0, 'G/loss_q': 0.0, 'G/new_rf_q': 0.0, 'G/new_lbls_q': 0.0, 'G/new_gp_q': 0.0, 'E/kld': 0.0, 'E/mse': 0.0, 'E/loss': 0.0 } for step, (x_spt, y_spt, x_qry, y_qry) in enumerate(train_loader, 0): x_spt, y_spt, x_qry, y_qry = x_spt.to(device), y_spt.to( device), x_qry.to(device), y_qry.to(device) BatchSize, setsz, c_, h, w = x_spt.size() querysz = x_qry.size(1) d_losses_q = [0.0 for _ in range(args.update_step)] g_losses_q = [0.0 for _ in range(args.update_step)] y_onehot = torch.cuda.FloatTensor(setsz, args.num_class) y_onehot_q = torch.cuda.FloatTensor(querysz, args.num_class) y_onehot_pre = torch.cuda.FloatTensor(setsz, args.num_class) for i in range(args.BatchSize): # This is inner loop not task inputs = Variable(x_spt[i]) labels = y_spt[i] real_feat = model(inputs) z = torch.Tensor( np.random.normal(0, 1, (setsz, args.latent_dim))).to(device) labels_q = y_qry[i] real_feat_q = model(x_qry[i]) z_q = torch.Tensor( np.random.normal( 0, 1, (querysz, args.latent_dim))).to(device) y_onehot.zero_() y_onehot.scatter(1, labels[:, None], 1) syn_label = y_onehot.to(device) y_onehot_q.zero_() y_onehot_q.scatter(1, labels_q[:, None], 1) syn_label_q = y_onehot_q.to(device) ############################# Train MetaGAN ########################### for k in range(args.update_step): fake_feat = generator(z, syn_label) fake_validity, disc_fake_acgan = discriminator( fake_feat, syn_label) real_validity, disc_real_acgan = discriminator( real_feat, syn_label) if current_task == 0: loss_aug = 0 * torch.sum(fake_validity) else: ind = list(range(len(pre_index))) embed_label_sythesis = [] for _ in range(setsz): np.random.shuffle(ind) embed_label_sythesis.append(pre_index[ind[0]]) embed_label_sythesis = np.asarray( embed_label_sythesis) embed_label_sythesis = torch.from_numpy( embed_label_sythesis) y_onehot_pre.zero_() y_onehot_pre.scatter( 1, embed_label_sythesis[:, None].to(device), 1) syn_label_pre = y_onehot_pre.to(device) pre_feat = generator(z, syn_label_pre) pre_feat_old = generator_old(z, syn_label_pre) loss_aug = loss_mse(pre_feat, pre_feat_old) # Adversarial loss (wasserstein) g_loss_lbls = criterion_softmax( disc_fake_acgan, labels.to(device)) d_loss_rf = -torch.mean(real_validity) + torch.mean( fake_validity) d_gradient_penalty = compute_gradient_penalty( discriminator, real_feat, fake_feat, syn_label).mean() d_loss_lbls = criterion_softmax( disc_real_acgan, labels.to(device)) d_loss = d_loss_rf + lambda_gp * d_gradient_penalty + 0.5 * ( d_loss_lbls + g_loss_lbls) g_loss_rf = -torch.mean(fake_validity) g_loss = g_loss_rf + lambda_lwf * old_task_factor * loss_aug + g_loss_lbls grad_d = torch.autograd.grad( d_loss, discriminator.parameters(), create_graph=True, retain_graph=True) grad_g = torch.autograd.grad(g_loss, generator.parameters(), create_graph=True, retain_graph=True) grad_d = clip_grad_by_norm_(grad_d, max_norm=5.0) grad_g = clip_grad_by_norm_(grad_g, max_norm=5) g_lr, d_lr = learned_lrs[k] for e, param in enumerate(discriminator.parameters(), 0): param.data = param.data.clone() - d_lr[e] * grad_d[ e] # args.update_lr * grad_d[e] for e, param in enumerate(generator.parameters(), 0): param.data = param.data.clone() - g_lr[e] * grad_g[ e] # args.update_lr * grad_g[e] fake_feat_q = generator(z_q, syn_label_q) fake_validity_q, disc_fake_acgan_q = discriminator( fake_feat_q, syn_label_q) real_validity_q, disc_real_acgan_q = discriminator( real_feat_q, syn_label_q) # Adversarial loss query d_loss_rf_q = -torch.mean( real_validity_q) + torch.mean(fake_validity_q) d_gradient_penalty_q = compute_gradient_penalty( discriminator, real_feat_q, fake_feat_q, syn_label_q).mean() d_loss_lbls_q = criterion_softmax( disc_real_acgan_q, labels_q.to(device)) d_loss_q = d_loss_rf_q + lambda_gp * d_gradient_penalty_q + d_loss_lbls_q d_losses_q[ k] = d_losses_q[k] + d_loss_q # + d_loss_lbls_q g_loss_rf_q = -torch.mean(fake_validity_q) g_loss_lbls_q = criterion_softmax( disc_fake_acgan_q, labels_q.to(device)) g_loss_q = g_loss_rf_q + g_loss_lbls_q # + lambda_lwf*old_task_factor * loss_aug_q g_losses_q[k] = g_losses_q[k] + g_loss_lbls_q #with torch.autograd.detect_anomaly(): optimizer_D.zero_grad() optimizer_G.zero_grad() optimizer_lr.zero_grad() d_loss_q_total = d_losses_q[-1].clone() / args.BatchSize g_loss_q_total = g_losses_q[-1].clone() / args.BatchSize d_loss_q_total.backward() g_loss_q_total.backward() torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 5) torch.nn.utils.clip_grad_norm_(generator.parameters(), 5) optimizer_D.step() optimizer_G.step() optimizer_lr.step() scheduler_G.step() scheduler_G.step() loss_log['D/loss'] += d_loss.item() loss_log['D/new_rf'] += d_loss_rf.item() loss_log['D/new_lbls'] += d_loss_lbls.item() #!!! loss_log['D/new_gp'] += d_gradient_penalty.item( ) if lambda_gp != 0 else 0 loss_log['D/loss_q'] += d_loss_q_total.item() #loss_log['D/new_rf_q'] += d_loss_rf_q.item() #loss_log['D/new_lbls_q'] += d_loss_lbls_q.item() #!!! #loss_log['D/new_gp_q'] += d_gradient_penalty_q.item() if lambda_gp != 0 else 0 del d_loss_rf, d_loss_lbls loss_log['G/loss'] += g_loss.item() loss_log['G/new_rf'] += g_loss_rf.item() loss_log['G/new_lbls'] += g_loss_lbls.item() #! loss_log['G/loss_q'] += g_loss_q_total.item() #loss_log['G/new_rf_q'] += g_loss_rf_q.item() #loss_log['G/new_lbls_q'] += g_loss_lbls_q.item() #!!! #loss_log['G/new_classifier'] += 0 #! loss_log['G/prev_mse'] += loss_aug.item( ) if lambda_lwf != 0 else 0 del g_loss_rf, g_loss_lbls print( '[GAN Epoch %05d]\t D Total Loss: %.3f \t G Total Loss: %.3f \t LwF Loss: %.3f' % (epoch + 1, loss_log['D/loss'], loss_log['G/loss'], loss_log['G/prev_rf'])) print( '[GAN Epoch %05d]\t D Total Loss Query: %.3f \t G Total Loss Query: %.3f \t' % (epoch + 1, loss_log['D/loss_q'], loss_log['G/loss_q'])) for k, v in loss_log.items(): if v != 0: tb_writer.add_scalar( 'Task {} - GAN/{}'.format(current_task, k), v, epoch + 1) if epoch == args.epochs_gan - 1: torch.save( generator, os.path.join( log_dir, 'task_' + str(current_task).zfill(2) + '_%d_model_generator.pkl' % epoch)) torch.save( discriminator, os.path.join( log_dir, 'task_' + str(current_task).zfill(2) + '_%d_model_discriminator.pkl' % epoch)) tb_writer.close() prototype = compute_prototype(model, train_loader, batch_size=args.BatchSize) #! return prototype
def main(args): # s_ = time.time() save_dir = args.save_dir mkdir_if_missing(save_dir) sys.stdout = logging.Logger(os.path.join(save_dir, 'log.txt')) display(args) start = 0 model = models.create(args.net, pretrained=True, dim=args.dim) # for vgg and densenet if args.resume is None: model_dict = model.state_dict() else: # resume model print('load model from {}'.format(args.resume)) chk_pt = load_checkpoint(args.resume) weight = chk_pt['state_dict'] start = chk_pt['epoch'] model.load_state_dict(weight) model = torch.nn.DataParallel(model) model = model.cuda() # freeze BN if args.freeze_BN is True: print(40 * '#', '\n BatchNorm frozen') model.apply(set_bn_eval) else: print(40 * '#', 'BatchNorm NOT frozen') # Fine-tune the model: the learning rate for pre-trained parameter is 1/10 new_param_ids = set(map(id, model.module.classifier.parameters())) new_params = [ p for p in model.module.parameters() if id(p) in new_param_ids ] base_params = [ p for p in model.module.parameters() if id(p) not in new_param_ids ] param_groups = [{ 'params': base_params, 'lr_mult': 0.0 }, { 'params': new_params, 'lr_mult': 1.0 }] print('initial model is save at %s' % save_dir) optimizer = torch.optim.Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay) criterion = losses.create(args.loss, margin=args.margin, alpha=args.alpha, base=args.loss_base).cuda() # Decor_loss = losses.create('decor').cuda() data = DataSet.create(args.data, ratio=args.ratio, width=args.width, origin_width=args.origin_width, root=args.data_root) train_loader = torch.utils.data.DataLoader( data.train, batch_size=args.batch_size, sampler=FastRandomIdentitySampler(data.train, num_instances=args.num_instances), drop_last=True, pin_memory=True, num_workers=args.nThreads) # save the train information for epoch in range(start, args.epochs): train(epoch=epoch, model=model, criterion=criterion, optimizer=optimizer, train_loader=train_loader, args=args) if epoch == 1: optimizer.param_groups[0]['lr_mul'] = 0.1 if (epoch + 1) % args.save_step == 0 or epoch == 0: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': (epoch + 1), }, is_best=False, fpath=osp.join( args.save_dir, 'ckp_ep' + str(epoch + 1) + '.pth.tar'))
if __name__ == '__main__': if len(sys.argv) != 2: print("Usage: python main.py result_sha(e.g., car_3d_det_test)") sys.exit(1) result_sha = sys.argv[1] save_root = '/home/eshan/AB3DMOT/results' det_id2str = {1: 'Pedestrian', 2: 'Car', 3: 'Cyclist'} seq_file_list, num_seq = load_list_from_folder( os.path.join('/home/eshan/AB3DMOT/data/KITTI', result_sha)) total_time = 0.0 total_frames = 0 save_dir = os.path.join(save_root, result_sha) mkdir_if_missing(save_dir) eval_dir = os.path.join(save_dir, 'data') mkdir_if_missing(eval_dir) for seq_file in seq_file_list: _, seq_name, _ = fileparts(seq_file) mot_tracker = AB3DMOT() seq_dets = np.loadtxt(seq_file, delimiter=',') # load detections eval_file = os.path.join(eval_dir, seq_name + '.txt') eval_file = open(eval_file, 'w') save_trk_dir = os.path.join(save_dir, 'trk_withid', seq_name) mkdir_if_missing(save_trk_dir) print("Processing %s." % (seq_name)) for frame in range(int(seq_dets[:, 0].min()), int(seq_dets[:, 0].max()) + 1): save_trk_file = os.path.join(save_trk_dir, '%06d.txt' % frame) save_trk_file = open(save_trk_file, 'w')
def prepare_student_data(nb_teachers, save=False): """ Takes a dataset name and the size of the teacher ensemble and prepares training data for the student model, according to parameters indicated in flags above. :param dataset: string corresponding to mnist, cifar10, or svhn :param nb_teachers: number of teachers (in the ensemble) to learn from :param save: if set to True, will dump student training labels predicted by the ensemble of teachers (with Laplacian noise) as npy files. It also dumps the clean votes for each class (without noise) and the labels assigned by teachers :return: pairs of (data, labels) to be used for student training and testing """ # Load the dataset if config.dataset == 'celeba': dataset = data_manager.init_img_dataset(root=config.data_dir, name=config.dataset) test_data = dataset.test_data test_labels = dataset.test_label train_data = dataset.train_data train_labels = dataset.train_label elif config.dataset =='market': data_dir = '../dataset/market1501' train_dataset = Train_Dataset(data_dir, dataset_name=dataset_dict[config.dataset], train_val='train') test_dataset = Test_Dataset(data_dir, dataset_name=dataset_dict[config.dataset], query_gallery='gallery') train_data = train_dataset.train_data train_labels = train_dataset.train_label test_data = test_dataset.data test_labels = test_dataset.label train_labels = np.array(train_labels,dtype =np.int32) test_labels = np.array(test_labels,dtype = np.int32) print('len of total test data in market',len(test_labels)) else: return False # Make sure there is data leftover to be used as a test set assert config.stdnt_share < len(test_data) ori_test_data = test_data # for test train_data, test_data = extract_feature(train_data, test_data) stdnt_data = test_data[:config.stdnt_share] # the remaining 1000 records is the holdout for evaluating share_index =np.random.choice(test_data[:-1000].shape[0],config.stdnt_share) stdnt_data = test_data[share_index] picked_stdnt_data = [ori_test_data[idx] for idx in share_index] num_train = train_data.shape[0] teachers_preds = np.zeros([stdnt_data.shape[0], config.nb_labels]) tau_teachers_preds=[] # a weighted teacher predtion with clippling for idx in range(len(stdnt_data)): if idx % 100 == 0: print('idx=', idx) query_data = stdnt_data[idx] select_teacher = np.random.choice(train_data.shape[0], int(prob * num_train)) dis = np.linalg.norm(train_data[select_teacher] - query_data, axis=1) k_index = select_teacher[np.argsort(dis)[:config.nb_teachers]] # sum over the number of teachers, which make it easy to compute their votings if config.use_tau: tau_teachers_preds.append(tau_limit(train_labels[k_index,:])) teachers_preds[idx] = np.sum(train_labels[k_index, :], axis=0) teachers_preds = np.asarray(teachers_preds, dtype=np.int32) if config.use_tau: preds_tau = np.asarray(tau_teachers_preds, dtype = np.float32) acct.compose_poisson_subsampled_mechanisms(gaussian, prob, coeff=config.stdnt_share) count_zero_list = config.nb_teachers * np.ones([config.stdnt_share,config.nb_labels]) - teachers_preds idx, stdnt_labels = aggregation.aggregation_knn(teachers_preds, config.gau_scale,count_zero_list=count_zero_list) else: acct.compose_poisson_subsampled_mechanisms(gaussian, prob, coeff=config.stdnt_share) idx, stdnt_labels = aggregation.aggregation_knn(teachers_preds, config.gau_scale) # compute privacy loss print("Composition of student subsampled Gaussian mechanisms gives ", (acct.get_eps(delta), delta)) # Print accuracy of aggregated label #ac_ag_labels = hamming_accuracy(stdnt_labels, test_labels[:config.stdnt_share], torch=False) ac_ag_labels = hamming_accuracy(stdnt_labels, test_labels[share_index], torch=False) precision = hamming_precision(stdnt_labels, test_labels[share_index], torch=False) print("Accuracy of the aggregated labels: " + str(ac_ag_labels)) print('precision of the aggregated labels'+str(precision)) current_eps = acct.get_eps(config.delta) # Store unused part of test set for use as a test set after student training stdnt_test_data = ori_test_data[-1000:] stdnt_test_labels = test_labels[-1000:] if save: # Prepare filepath for numpy dump of labels produced by noisy aggregation dir_path = os.path.join(config.save_model, 'knn_num_neighbor_' + str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) filepath = dir_path + '_knn_voting.npy' #NOLINT(long-line) # Dump student noisy labels array with open(filepath, 'wb') as file_obj: np.save(file_obj, teachers_preds) return picked_stdnt_data, stdnt_labels, stdnt_test_data, stdnt_test_labels
def dataset_mpc(dataset, imset, seed=1, n=100, out_dir=None, crop_size=80): # Load the building d = sbpd.get_dataset(dataset, imset) ns = d.get_imset() camera_param = utils.Foo(width=225, height=225, z_near=0.01, z_far=20.0, fov_horizontal=60., fov_vertical=60., modalities=['rgb'], img_channels=3, im_resize=1.) r_obj = sr.get_r_obj(camera_param) for name in ns: for flip in [True, False]: tt = d.load_data(name, flip) tt.set_r_obj(r_obj) tt.load_building_into_scene() traversible_cc, resolution = tt.traversible, tt.map.resolution traversible_map = tt.map.traversible * 1. starts, ends = mu.sample_positions_on_map(seed, traversible_cc, resolution, n) # Crop out map from this location. loc = starts * 1. diff = ends - starts theta = np.arctan2(diff[:, 1], diff[:, 0])[:, np.newaxis] x_axis = np.concatenate([np.cos(theta), np.sin(theta)], axis=1) y_axis = np.concatenate( [np.cos(theta + np.pi / 2.), np.sin(theta + np.pi / 2.)], axis=1) crops = mu.generate_egocentric_maps([traversible_map], [1.0], [crop_size], starts, x_axis, y_axis, dst_theta=np.pi / 2.0) # Render out images for each location. nodes = np.concatenate([loc, theta / tt.robot.delta_theta], axis=1) imgs = tt.render_nodes(nodes) r_obj.clear_scene() # Write the maps and the images in a directory. if out_dir is not None: out_dir_ = os.path.join( out_dir, '{:d}'.format(seed), '{:s}_{:s}_{:d}'.format(dataset, name, flip)) logging.error('out_dir: %s', out_dir_) utils.mkdir_if_missing(out_dir_) for i in range(n): file_name = os.path.join(out_dir_, 'img_{:06d}.jpg'.format(i)) cv2.imwrite(file_name, imgs[i][:, :, ::-1].astype(np.uint8)) file_name = os.path.join(out_dir_, 'map_{:06d}.png'.format(i)) cv2.imwrite(file_name, (crops[0][i, :, :] * 255).astype(np.uint8)) out_file_name = os.path.join( out_dir, 'maps', '{:s}_{:s}_{:d}.pkl'.format(dataset, name, flip)) vv = vars(tt.map) vv['traversible_cc'] = tt.traversible # save_variables(out_file_name, list(vv.values()), list(vv.keys()), True) output_file = os.path.join(out_dir_, 'vis.html') create_webpage(output_file, n)
def main(args): s_ = time.time() # 训练日志保存 log_dir = args.log_dir mkdir_if_missing(log_dir) sys.stdout = logging.Logger(os.path.join(log_dir, 'log.txt')) display(args) if args.r is None: model = models.create(args.net) model = load_parameter(model) else: # resume model print('Resume from model at Epoch %d' % args.start) model = torch.load(args.r) model = model.cuda() torch.save(model, os.path.join(log_dir, 'model.pkl')) print('initial model is save at %s' % log_dir) # fine tune the model: the learning rate for pre-trained parameter is 1/10 new_param_ids = set.union( set(map(id, model.Embedding.parameters())), set(map(id, model.attention_blocks.parameters()))) new_params = [p for p in model.parameters() if id(p) in new_param_ids] base_params = [p for p in model.parameters() if id(p) not in new_param_ids] param_groups = [{ 'params': base_params, 'lr_mult': 0.0 }, { 'params': new_params, 'lr_mult': 1.0 }] optimizer = torch.optim.Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay) if args.loss == 'bin': criterion = losses.create(args.loss, margin=args.margin, alpha=args.alpha).cuda() Div = losses.create('div').cuda() else: criterion = losses.create(args.loss).cuda() data = DataSet.create(args.data, root=None) train_loader = torch.utils.data.DataLoader( data.train, batch_size=args.BatchSize, sampler=RandomIdentitySampler(data.train, num_instances=args.num_instances), drop_last=True, num_workers=args.nThreads) # save the train information epoch_list = list() loss_list = list() pos_list = list() neg_list = list() for epoch in range(args.start, args.epochs): epoch_list.append(epoch) running_loss = 0.0 divergence = 0.0 running_pos = 0.0 running_neg = 0.0 if epoch == 2: param_groups[0]['lr_mult'] = 0.1 for i, data in enumerate(train_loader, 0): inputs, labels = data # wrap them in Variable inputs = Variable(inputs.cuda()) # type of labels is Variable cuda.Longtensor labels = Variable(labels).cuda() optimizer.zero_grad() embed_feat = model(inputs) loss, inter_, dist_ap, dist_an = criterion(embed_feat, labels) div = Div(embed_feat) loss_ = loss + args.theta * div if not type(loss) == torch.Tensor: print('One time con not back-ward') continue loss_.backward() optimizer.step() running_loss += loss.item() divergence += div.item() running_neg += dist_an running_pos += dist_ap if epoch == 0 and i == 0: print(50 * '#') print('Train Begin -- HA-HA-HA-HA-AH-AH-AH-AH --') loss_list.append(running_loss) pos_list.append(running_pos / i) neg_list.append(running_neg / i) print( '[Epoch %05d]\t Loss: %.2f \t Divergence: %.2f \t Accuracy: %.2f \t Pos-Dist: %.2f \t Neg-Dist: %.2f' % (epoch + 1, running_loss, divergence, inter_, dist_ap, dist_an)) if epoch % args.save_step == 0: torch.save(model, os.path.join(log_dir, '%d_model.pkl' % epoch)) np.savez(os.path.join(log_dir, "result.npz"), epoch=epoch_list, loss=loss_list, pos=pos_list, neg=neg_list) t = time.time() - s_ print('training takes %.2f hour' % (t / 3600))
parser.add_argument('-orth_cof', type=float, default=0e-3, help='try to make the last linear weight matrix to ' 'approximate the orthogonal matrix') args = parser.parse_args() print(args.nums) print(type(args.nums)) if args.log_dir is None: log_dir = os.path.join('checkpoints', args.loss) else: log_dir = os.path.join('checkpoints', args.log_dir) mkdir_if_missing(log_dir) # write log sys.stdout = logging.Logger(os.path.join(log_dir, 'log.txt')) # display information of current training print('train on dataset %s' % args.data) print('batch size is: %d' % args.BatchSize) print('num_instance is %d' % args.num_instances) print('dimension of the embedding space is %d' % args.dim) print('log dir is: %s' % args.log_dir) print('the network is : %s' % args.net) print('loss function for training is: %s' % args.loss) print('learn rate : %f' % args.lr) print('base parameter de learn rate is : %f' % args.base) print('the orthogonal weight regular is %f ' % args.orth_cof)
def prepare_student_data(nb_teachers, save=False): """ Takes a dataset name and the size of the teacher ensemble and prepares training data for the student model, according to parameters indicated in flags above. :param dataset: string corresponding to mnist, cifar10, or svhn :param nb_teachers: number of teachers (in the ensemble) to learn from :param save: if set to True, will dump student training labels predicted by the ensemble of teachers (with Laplacian noise) as npy files. It also dumps the clean votes for each class (without noise) and the labels assigned by teachers :return: pairs of (data, labels) to be used for student training and testing """ # Load the dataset if config.dataset == 'celeba': dataset = data_manager.init_img_dataset(root=config.data_dir, name=config.dataset) test_data = dataset.test_data test_labels = dataset.test_label elif config.dataset == 'market': data_dir = '../dataset/market1501' test_dataset = Test_Dataset(data_dir, dataset_name=dataset_dict[config.dataset], query_gallery='gallery') test_data = test_dataset.data test_labels = test_dataset.label test_labels = np.array(test_labels, dtype=np.int32) else: print("Check value of dataset flag") return False # Make sure there is data leftover to be used as a test set assert config.stdnt_share < len(test_data) # Prepare [unlabeled] student training data (subset of test set) stdnt_data = test_data[:config.stdnt_share] # Compute teacher predictions for student training data if config.reuse_vote: #reuse previous saved clean votes, but stdnt_share maybe various #dir_path = os.path.join(config.save_model,'pate_'+str(config.nb_teachers)) dir_path = os.path.join(config.save_model, config.dataset) dir_path = os.path.join(dir_path, 'pate_num_teacher_' + str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) filepath = dir_path + '/_teacher_votes.npy' # Prepare filepath for numpy dump of labels produced by noisy aggregation teachers_preds = np.load(filepath) teachers_preds = teachers_preds[:config.stdnt_share] ori_filepath = dir_path + '_ori_teacher_votes.npy' ori_teachers_preds = np.load(ori_filepath) else: teachers_preds = ensemble_preds(nb_teachers, stdnt_data) ori_teachers_preds = teachers_preds # in the shape of (nb_teacher, nb_data, dim) teachers_preds = np.sum(teachers_preds, axis=0) dir_path = os.path.join(config.save_model, config.dataset) dir_path = os.path.join(dir_path, 'pate_num_teacher_' + str(config.nb_teachers)) utils.mkdir_if_missing(dir_path) filepath = dir_path + '/_teacher_votes.npy' ori_filepath = dir_path + '_ori_teacher_votes.npy' with open(filepath, mode='wb') as file_obj: np.save(file_obj, teachers_preds) with open(ori_filepath, mode='wb') as file_obj: np.save(file_obj, ori_teachers_preds) if config.use_tau: tau_teachers_preds = np.zeros(teachers_preds.shape) for idx in range(len(tau_teachers_preds)): tau_teachers_preds[idx] = tau_limit(ori_teachers_preds[:, idx, :]) preds_tau = np.asarray(tau_teachers_preds, dtype=np.float32) print('preds_tau', preds_tau[1, ]) count_zero_list = config.nb_teachers * np.ones( [config.stdnt_share, config.nb_labels]) - teachers_preds print('shape of count_zero', count_zero_list.shape) idx, stdnt_labels = aggregation.aggregation_knn( teachers_preds, config.gau_scale, count_zero_list=count_zero_list) acct.compose_mechanism(gaussian, coeff=config.stdnt_share) else: acct.compose_mechanism(gaussian, coeff=config.stdnt_share) idx, stdnt_labels = aggregation.aggregation_knn( teachers_preds, config.gau_scale) print('shape of teachers_pred', teachers_preds.shape) # Aggregate teacher predictions to get student training labels # Print accuracy of aggregated label ac_ag_labels = hamming_accuracy(stdnt_labels, test_labels[:config.stdnt_share], torch=False) print("Accuracy of the aggregated labels: " + str(ac_ag_labels)) current_eps = acct.get_eps(config.delta) print('eps after data independent composition', current_eps) # Store unused part of test set for use as a test set after student training stdnt_test_data = test_data[config.stdnt_share:] stdnt_test_labels = test_labels[config.stdnt_share:] return stdnt_data, stdnt_labels, stdnt_test_data, stdnt_test_labels
'/Users/idofarhi/Documents/Thesis/Data/frames/' + folder_name + '/4CH' } # Test data: # source_directory = {'apex': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/raw/apex', # 'mitral': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/raw/mitral', # 'papillary': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/raw/papillary'} # # target_directory = {'apex': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/' + folder_name + '/apex', # 'mitral': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/' + folder_name + '/mitral', # 'papillary': '/Users/idofarhi/Documents/Thesis/Data/test_set/frames/' + folder_name + '/papillary'} ################################################# for folder in target_directory.values(): mkdir_if_missing(folder) video_list = {} file_names = get_frame_sequences(source_directory['apex'][:-5], class_folders=view_list) for view in view_list: video_list[view] = list(file_names[view].keys()) # for each view... for view in view_list: # get a list of all files in the relevant view directory (file list) file_list = get_files_list(source_directory[view]) print('Running view:', view) # run through each video... for video in tqdm(video_list[view]):
def main(args): # 训练日志保存 log_dir = os.path.join('checkpoints', args.log_dir) mkdir_if_missing(log_dir) sys.stdout = logging.Logger(os.path.join(log_dir, 'log.txt')) display(args) if args.r is None: model = models.create(args.net, Embed_dim=args.dim) # load part of the model model_dict = model.state_dict() # print(model_dict) if args.net == 'bn': pretrained_dict = torch.load( 'pretrained_models/bn_inception-239d2248.pth') else: pretrained_dict = torch.load( 'pretrained_models/inception_v3_google-1a9a5a14.pth') pretrained_dict = { k: v for k, v in pretrained_dict.items() if k in model_dict } model_dict.update(pretrained_dict) model.load_state_dict(model_dict) else: # resume model model = torch.load(args.r) model = model.cuda() torch.save(model, os.path.join(log_dir, 'model.pkl')) print('initial model is save at %s' % log_dir) # fine tune the model: the learning rate for pre-trained parameter is 1/10 new_param_ids = set(map(id, model.Embed.parameters())) new_params = [p for p in model.parameters() if id(p) in new_param_ids] base_params = [p for p in model.parameters() if id(p) not in new_param_ids] param_groups = [{ 'params': base_params, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] optimizer = torch.optim.Adam(param_groups, lr=args.lr, weight_decay=args.weight_decay) criterion = losses.create(args.loss, alpha=args.alpha, k=args.k).cuda() data = DataSet.create(args.data, root=None, test=False) train_loader = torch.utils.data.DataLoader( data.train, batch_size=args.BatchSize, sampler=RandomIdentitySampler(data.train, num_instances=args.num_instances), drop_last=True, num_workers=args.nThreads) for epoch in range(args.start, args.epochs): running_loss = 0.0 for i, data in enumerate(train_loader, 0): inputs, labels = data # wrap them in Variable inputs = Variable(inputs.cuda()) labels = Variable(labels).cuda() optimizer.zero_grad() embed_feat = model(inputs) loss, inter_, dist_ap, dist_an = criterion(embed_feat, labels) if args.orth > 0: loss = orth_reg(model, loss, cof=args.orth) loss.backward() optimizer.step() running_loss += loss.data[0] if epoch == 0 and i == 0: print(50 * '#') print('Train Begin -- HA-HA-HA') print( '[Epoch %05d]\t Loss: %.3f \t Accuracy: %.3f \t Pos-Dist: %.3f \t Neg-Dist: %.3f' % (epoch + 1, running_loss, inter_, dist_ap, dist_an)) if epoch % args.save_step == 0: torch.save(model, os.path.join(log_dir, '%d_model.pkl' % epoch))
def _preprocess(self): """ This function is a bit complex and ugly, what it does is 1. Extract data from cuhk-03.mat and save as png images. 2. Create 20 classic splits. (Li et al. CVPR'14) 3. Create new split. (Zhong et al. CVPR'17) """ print("Note: if root path is changed, the previously generated json files need to be re-generated (delete them first)") if osp.exists(self.imgs_labeled_dir) and \ osp.exists(self.imgs_detected_dir) and \ osp.exists(self.split_classic_det_json_path) and \ osp.exists(self.split_classic_lab_json_path) and \ osp.exists(self.split_new_det_json_path) and \ osp.exists(self.split_new_lab_json_path): return mkdir_if_missing(self.imgs_detected_dir) mkdir_if_missing(self.imgs_labeled_dir) print("Extract image data from {} and save as png".format(self.raw_mat_path)) mat = h5py.File(self.raw_mat_path, 'r') def _deref(ref): return mat[ref][:].T def _process_images(img_refs, campid, pid, save_dir): img_paths = [] # Note: some persons only have images for one view for imgid, img_ref in enumerate(img_refs): img = _deref(img_ref) # skip empty cell if img.size == 0 or img.ndim < 3: continue # images are saved with the following format, index-1 (ensure uniqueness) # campid: index of camera pair (1-5) # pid: index of person in 'campid'-th camera pair # viewid: index of view, {1, 2} # imgid: index of image, (1-10) viewid = 1 if imgid < 5 else 2 img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(campid+1, pid+1, viewid, imgid+1) img_path = osp.join(save_dir, img_name) imsave(img_path, img) img_paths.append(img_path) return img_paths def _extract_img(name): print("Processing {} images (extract and save) ...".format(name)) meta_data = [] imgs_dir = self.imgs_detected_dir if name == 'detected' else self.imgs_labeled_dir for campid, camp_ref in enumerate(mat[name][0]): camp = _deref(camp_ref) num_pids = camp.shape[0] for pid in range(num_pids): img_paths = _process_images(camp[pid,:], campid, pid, imgs_dir) assert len(img_paths) > 0, "campid{}-pid{} has no images".format(campid, pid) meta_data.append((campid+1, pid+1, img_paths)) print("done camera pair {} with {} identities".format(campid+1, num_pids)) return meta_data meta_detected = _extract_img('detected') meta_labeled = _extract_img('labeled') def _extract_classic_split(meta_data, test_split): train, test = [], [] num_train_pids, num_test_pids = 0, 0 num_train_imgs, num_test_imgs = 0, 0 for i, (campid, pid, img_paths) in enumerate(meta_data): if [campid, pid] in test_split: for img_path in img_paths: camid = int(osp.basename(img_path).split('_')[2]) test.append((img_path, num_test_pids, camid)) num_test_pids += 1 num_test_imgs += len(img_paths) else: for img_path in img_paths: camid = int(osp.basename(img_path).split('_')[2]) train.append((img_path, num_train_pids, camid)) num_train_pids += 1 num_train_imgs += len(img_paths) return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs print("Creating classic splits (# = 20) ...") splits_classic_det, splits_classic_lab = [], [] for split_ref in mat['testsets'][0]: test_split = _deref(split_ref).tolist() # create split for detected images train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \ _extract_classic_split(meta_detected, test_split) splits_classic_det.append({ 'train': train, 'query': test, 'gallery': test, 'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs, 'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs, 'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs, }) # create split for labeled images train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \ _extract_classic_split(meta_labeled, test_split) splits_classic_lab.append({ 'train': train, 'query': test, 'gallery': test, 'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs, 'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs, 'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs, }) write_json(splits_classic_det, self.split_classic_det_json_path) write_json(splits_classic_lab, self.split_classic_lab_json_path) def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel): tmp_set = [] unique_pids = set() for idx in idxs: img_name = filelist[idx][0] camid = int(img_name.split('_')[2]) pid = pids[idx] if relabel: pid = pid2label[pid] img_path = osp.join(img_dir, img_name) tmp_set.append((img_path, int(pid), camid)) unique_pids.add(pid) return tmp_set, len(unique_pids), len(idxs) def _extract_new_split(split_dict, img_dir): train_idxs = split_dict['train_idx'].flatten() - 1 # index-0 pids = split_dict['labels'].flatten() train_pids = set(pids[train_idxs]) pid2label = {pid: label for label, pid in enumerate(train_pids)} query_idxs = split_dict['query_idx'].flatten() - 1 gallery_idxs = split_dict['gallery_idx'].flatten() - 1 filelist = split_dict['filelist'].flatten() train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True) query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False) gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False) return train_info, query_info, gallery_info print("Creating new splits for detected images (767/700) ...") train_info, query_info, gallery_info = _extract_new_split( loadmat(self.split_new_det_mat_path), self.imgs_detected_dir, ) splits = [{ 'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0], 'num_train_pids': train_info[1], 'num_train_imgs': train_info[2], 'num_query_pids': query_info[1], 'num_query_imgs': query_info[2], 'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2], }] write_json(splits, self.split_new_det_json_path) print("Creating new splits for labeled images (767/700) ...") train_info, query_info, gallery_info = _extract_new_split( loadmat(self.split_new_lab_mat_path), self.imgs_labeled_dir, ) splits = [{ 'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0], 'num_train_pids': train_info[1], 'num_train_imgs': train_info[2], 'num_query_pids': query_info[1], 'num_query_imgs': query_info[2], 'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2], }] write_json(splits, self.split_new_lab_json_path)
def track_nuscenes(data_root: str, detection_file: str, save_dir: str, eval_set: str = 'val', covariance_id: int = 0, match_distance: str = 'iou', match_threshold: float = 0.1, match_algorithm: str = 'h', use_angular_velocity: bool = False): ''' submission { "meta": { "use_camera": <bool> -- Whether this submission uses camera data as an input. "use_lidar": <bool> -- Whether this submission uses lidar data as an input. "use_radar": <bool> -- Whether this submission uses radar data as an input. "use_map": <bool> -- Whether this submission uses map data as an input. "use_external": <bool> -- Whether this submission uses external data as an input. }, "results": { sample_token <str>: List[sample_result] -- Maps each sample_token to a list of sample_results. } } ''' if 'train' in eval_set: version = 'v1.0-trainval' elif 'val' in eval_set: version = 'v1.0-trainval' elif 'mini' in eval_set: version = 'v1.0-mini' elif 'test' in eval_set: version = 'v1.0-test' else: version = eval_set print("WARNING: Unknown subset version: '{}'".format(version)) nusc = NuScenes(version=version, dataroot=data_root, verbose=True) mkdir_if_missing(save_dir) output_path = os.path.join(save_dir, 'probabilistic_tracking_results.json') results = {} total_time = 0.0 total_frames = 0 with open(detection_file) as f: data = json.load(f) assert 'results' in data, 'Error: No field `results` in result file. Please note that the result format changed.' \ 'See https://www.nuscenes.org/object-detection for more information.' all_results = EvalBoxes.deserialize(data['results'], DetectionBox) meta = data['meta'] print('meta: ', meta) print("Loaded results from {}. Found detections for {} samples.".format( detection_file, len(all_results.sample_tokens))) processed_scene_tokens = set() for sample_token_idx in tqdm(range(len(all_results.sample_tokens))): sample_token = all_results.sample_tokens[sample_token_idx] scene_token = nusc.get('sample', sample_token)['scene_token'] if scene_token in processed_scene_tokens: continue first_sample_token = nusc.get('scene', scene_token)['first_sample_token'] current_sample_token = first_sample_token mot_trackers = { tracking_name: AB3DMOT(covariance_id, tracking_name=tracking_name, use_angular_velocity=use_angular_velocity, tracking_nuscenes=True) for tracking_name in NUSCENES_TRACKING_NAMES } while current_sample_token != '': results[current_sample_token] = [] dets = { tracking_name: [] for tracking_name in NUSCENES_TRACKING_NAMES } info = { tracking_name: [] for tracking_name in NUSCENES_TRACKING_NAMES } for box in all_results.boxes[current_sample_token]: if box.detection_name not in NUSCENES_TRACKING_NAMES: continue q = Quaternion(box.rotation) angle = q.angle if q.axis[2] > 0 else -q.angle #print('box.rotation, angle, axis: ', box.rotation, q.angle, q.axis) #print('box.rotation, angle, axis: ', q.angle, q.axis) #[h, w, l, x, y, z, rot_y] detection = np.array([ box.size[2], box.size[0], box.size[1], box.translation[0], box.translation[1], box.translation[2], angle ]) #print('detection: ', detection) information = np.array([box.detection_score]) dets[box.detection_name].append(detection) info[box.detection_name].append(information) dets_all = { tracking_name: { 'dets': np.array(dets[tracking_name]), 'info': np.array(info[tracking_name]) } for tracking_name in NUSCENES_TRACKING_NAMES } total_frames += 1 start_time = time.time() for tracking_name in NUSCENES_TRACKING_NAMES: if dets_all[tracking_name]['dets'].shape[0] > 0: trackers = mot_trackers[tracking_name].update( dets_all[tracking_name], match_distance, match_threshold, match_algorithm, scene_token) # (N, 9) # (h, w, l, x, y, z, rot_y), tracking_id, tracking_score # print('trackers: ', trackers) for i in range(trackers.shape[0]): sample_result = format_sample_result( current_sample_token, tracking_name, trackers[i]) results[current_sample_token].append(sample_result) cycle_time = time.time() - start_time total_time += cycle_time # get next frame and continue the while loop current_sample_token = nusc.get('sample', current_sample_token)['next'] # left while loop and mark this scene as processed processed_scene_tokens.add(scene_token) # finished tracking all scenes, write output data output_data = {'meta': meta, 'results': results} with open(output_path, 'w') as outfile: json.dump(output_data, outfile) print("Total Tracking took: %.3f for %d frames or %.1f FPS" % (total_time, total_frames, total_frames / total_time))
sys.exit(1) data_type = sys.argv[2] if data_type != 'vehicle' and data_type != 'pedestrian' and data_type != 'cyclist': print("Usage: python main.py data.bin vehicle") sys.exit(1) data_root = "./dataset" data_file_name = fileparts(sys.argv[1])[1] + fileparts(sys.argv[1])[2] result_root = './results' result_file_name = data_file_name[:-4] + "_preds.bin" data_file = os.path.join(data_root, data_file_name) mkdir_if_missing(result_root) dataset = metrics_pb2.Objects() with open(data_file, 'rb') as f: buf = f.read() dataset.ParseFromString(buf) contexts = [] if len(dataset.objects) != 0: context = dataset.objects[0].context_name start = 0 i = 1 for data in dataset.objects[1:]: if data.context_name != context: contexts.append((start, i))
def train_fun(args, train_loader, feat_loader, current_task, fisher={}, prototype={}): log_dir = args.log_dir mkdir_if_missing(log_dir) sys.stdout = logging.Logger(os.path.join(log_dir, 'log.txt')) display(args) model = models.create(args.net, Embed_dim=args.dim) # load part of the model if args.method == 'Independent' or current_task == 0: model_dict = model.state_dict() if args.net == 'resnet32': pickle.load = partial(pickle.load, encoding="latin1") pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1") pretrained_dict = torch.load( 'pretrained_models/Finetuning_0_task_0_200_model_task2_cifar100_seed1993.pkl', map_location=lambda storage, loc: storage, pickle_module=pickle) pretrained_dict = pretrained_dict.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items( ) if k in model_dict and 'fc' not in k} model_dict.update(pretrained_dict) model.load_state_dict(model_dict) elif args.net == 'resnet18' and args.data == 'imagenet_sub': pickle.load = partial(pickle.load, encoding="latin1") pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1") pretrained_dict = torch.load( 'pretrained_models/Finetuning_0_task_0_200_model_task2_imagenet_sub_seed1993.pkl', map_location=lambda storage, loc: storage, pickle_module=pickle) pretrained_dict = pretrained_dict.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items( ) if k in model_dict and 'fc' not in k} model_dict.update(pretrained_dict) model.load_state_dict(model_dict) else: print (' Oops! That was no valid models. ') if args.method != 'Independent' and current_task > 0: model = torch.load(os.path.join(log_dir, args.method + '_' + args.exp + '_task_' + str(current_task-1) + '_%d_model.pkl' % int(args.epochs-1))) model_old = deepcopy(model) model_old.eval() model_old = freeze_model(model_old) model = model.cuda() torch.save(model, os.path.join(log_dir, args.method + '_' + args.exp + '_task_' + str(current_task) + '_pre_model.pkl')) print('initial model is save at %s' % log_dir) # fine tune the model: the learning rate for pre-trained parameter is 1/10 new_param_ids = set(map(id, model.Embed.parameters())) new_params = [p for p in model.parameters() if id(p) in new_param_ids] base_params = [p for p in model.parameters() if id(p) not in new_param_ids] param_groups = [ {'params': base_params, 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] criterion = losses.create(args.loss, margin=args.margin, num_instances=args.num_instances).cuda() optimizer = torch.optim.Adam( param_groups, lr=args.lr, weight_decay=args.weight_decay) scheduler = StepLR(optimizer, step_size=200, gamma=0.1) if args.data == 'cifar100' or args.data == 'imagenet_sub': if current_task > 0: model.eval() for epoch in range(args.start, args.epochs): running_loss = 0.0 running_lwf = 0.0 scheduler.step() for i, data in enumerate(train_loader, 0): inputs, labels = data # wrap them in Variable inputs = Variable(inputs.cuda()) labels = Variable(labels).cuda() optimizer.zero_grad() _, embed_feat = model(inputs) if current_task == 0: loss_aug = 0*torch.sum(embed_feat) else: if args.method == 'Finetuning' or args.method == 'Independent': loss_aug = 0*torch.sum(embed_feat) elif args.method == 'LwF': _, embed_feat_old = model_old(inputs) loss_aug = args.tradeoff * \ torch.sum((embed_feat-embed_feat_old).pow(2))/2. elif args.method == 'EWC' or args.method == 'MAS': loss_aug = 0 for (name, param), (_, param_old) in zip(model.named_parameters(), model_old.named_parameters()): loss_aug += args.tradeoff * \ torch.sum(fisher[name]*(param_old-param).pow(2))/2. if args.loss == 'MSLoss': loss = criterion(embed_feat, labels) inter_ = 0 dist_ap = 0 dist_an = 0 else: loss, inter_, dist_ap, dist_an = criterion(embed_feat, labels) loss += loss_aug loss.backward() optimizer.step() running_loss += loss.data[0] running_lwf += loss_aug.data[0] if epoch == 0 and i == 0: print(50*'#') print('Train Begin -- HA-HA-HA') print('[Epoch %05d]\t Total Loss: %.3f \t LwF Loss: %.3f \t Accuracy: %.3f \t Pos-Dist: %.3f \t Neg-Dist: %.3f' % (epoch + 1, running_loss, running_lwf, inter_, dist_ap, dist_an)) if epoch % args.save_step == 0: torch.save(model, os.path.join(log_dir, args.method + '_' + args.exp + '_task_' + str(current_task) + '_%d_model.pkl' % epoch)) if args.method == 'EWC' or args.method == 'MAS': fisher = fisher_matrix_diag( model, criterion, train_loader, number_samples=500) return fisher