""" Extract features from pre-trained networks. The main procedures are finetune and extract features. Finetune: Given an Imagenet pretrained model (such as ResNet50), finetune it on a dataset (we call it source) Extractor: After fine-tune, extract features on the target domain using finetuned models on source This class supports most image models: Alexnet, Resnet(xx), VGG. Other text or digit models can be easily extended using this code, see models.py for details. """ import argparse import data_load import models import numpy as np import torch import torch.nn as nn import torch.optim as optim import torchvision import time import copy import os # Command setting parser = argparse.ArgumentParser(description='Finetune') parser.add_argument('--model_name', type=str, help='model name', default='resnet50') parser.add_argument('--batchsize', type=int, help='batch size', default=64) parser.add_argument('--gpu', type=int, help='cuda id', default=0) parser.add_argument('--source', type=str, default='amazon') parser.add_argument('--target', type=str, default='webcam') parser.add_argument('--num_class', type=int, default=12) parser.add_argument('--dataset_path', type=str, default='../../data/Office31/Original_images/') parser.add_argument('--epoch', type=int, help='Train epochs', default=100) parser.add_argument('--momentum', type=float, help='Momentum', default=.9) parser.add_argument('--lr', type=float, default=1e-4) parser.add_argument('--finetune', type=int, help='Needs finetune or not', default=1) parser.add_argument('--extract', type=int, help='Needs extract features or not', default=1) args = parser.parse_args() # Parameter setting DEVICE = torch.device('cuda:' + str(args.gpu) if torch.cuda.is_available() else 'cpu') BATCH_SIZE = {'src': int(args.batchsize), 'tar': int(args.batchsize)} def get_optimizer(model): learning_rate = args.lr param_group = [] param_group += [{'params': model.base_network.parameters(), 'lr': learning_rate}] param_group += [{'params': model.classifier_layer.parameters(), 'lr': learning_rate * 10}] optimizer = optim.SGD(param_group, momentum=args.momentum) return optimizer # Schedule learning rate according to DANN if you want to (while I think this equation is wierd therefore I did not use this one) def lr_schedule(optimizer, epoch): def lr_decay(LR, n_epoch, e): return LR / (1 + 10 * e / n_epoch) ** 0.75 for i in range(len(optimizer.param_groups)): if i < len(optimizer.param_groups) - 1: optimizer.param_groups[i]['lr'] = lr_decay( LEARNING_RATE, N_EPOCH, epoch) else: optimizer.param_groups[i]['lr'] = lr_decay( LEARNING_RATE, N_EPOCH, epoch) * 10 def finetune(model, dataloaders, optimizer, criterion, best_model_path, use_lr_schedule=False): N_EPOCH = args.epoch best_model_wts = copy.deepcopy(model.state_dict()) since = time.time() best_acc = 0.0 acc_hist = [] for epoch in range(1, N_EPOCH + 1): if use_lr_schedule: lr_schedule(optimizer, epoch) for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval() total_loss, correct = 0, 0 for inputs, labels in dataloaders[phase]: inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) loss = criterion(outputs, labels) preds = torch.max(outputs, 1)[1] if phase == 'train': loss.backward() optimizer.step() total_loss += loss.item() * inputs.size(0) correct += torch.sum(preds == labels.data) epoch_loss = total_loss / len(dataloaders[phase].dataset) epoch_acc = correct.double() / len(dataloaders[phase].dataset) acc_hist.append([epoch_loss, epoch_acc]) print('Epoch: [{:02d}/{:02d}]---{}, loss: {:.6f}, acc: {:.4f}'.format(epoch, N_EPOCH, phase, epoch_loss, epoch_acc)) if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) torch.save(model.state_dict( ), 'save_model/best_{}_{}-{}.pth'.format(args.model_name, args.source, epoch)) time_pass = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_pass // 60, time_pass % 60)) print('------Best acc: {}'.format(best_acc)) model.load_state_dict(best_model_wts) torch.save(model.state_dict(), best_model_path) print('Best model saved!') return model, best_acc, acc_hist # Extract features for given intermediate layers # Currently, this only works for ResNet since AlexNet and VGGNET only have features and classifiers modules. # You will need to manually define a function in the forward function to extract features # (by letting it return features and labels). # Please follow digit_deep_network.py for reference. class FeatureExtractor(nn.Module): def __init__(self, model, extracted_layers): super(FeatureExtractor, self).__init__() self.model = model._modules['module'] if type( model) == torch.nn.DataParallel else model self.extracted_layers = extracted_layers def forward(self, x): outputs = [] for name, module in self.model._modules.items(): if name is "fc": x = x.view(x.size(0), -1) x = module(x) if name in self.extracted_layers: outputs.append(x) return outputs def extract_feature(model, dataloader, save_path, load_from_disk=True, model_path=''): if load_from_disk: model = models.Network(base_net=args.model_name, n_class=args.num_class) model.load_state_dict(torch.load(model_path)) model = model.to(DEVICE) model.eval() correct = 0 fea_all = torch.zeros(1,1+model.base_network.output_num()).to(DEVICE) with torch.no_grad(): for inputs, labels in dataloader: inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) feas = model.get_features(inputs) labels = labels.view(labels.size(0), 1).float() x = torch.cat((feas, labels), dim=1) fea_all = torch.cat((fea_all, x), dim=0) outputs = model(inputs) preds = torch.max(outputs, 1)[1] correct += torch.sum(preds == labels.data.long()) test_acc = correct.double() / len(dataloader.dataset) fea_numpy = fea_all.cpu().numpy() np.savetxt(save_path, fea_numpy[1:], fmt='%.6f', delimiter=',') print('Test acc: %f' % test_acc) # You may want to classify with 1nn after getting features def classify_1nn(data_train, data_test): ''' Classification using 1NN Inputs: data_train, data_test: train and test csv file path Outputs: yprediction and accuracy ''' from sklearn.neighbors import KNeighborsClassifier from sklearn.metrics import accuracy_score from sklearn.preprocessing import StandardScaler data = {'src': np.loadtxt(data_train, delimiter=','), 'tar': np.loadtxt(data_test, delimiter=','), } Xs, Ys, Xt, Yt = data['src'][:, :-1], data['src'][:, - 1], data['tar'][:, :-1], data['tar'][:, -1] Xs = StandardScaler(with_mean=0, with_std=1).fit_transform(Xs) Xt = StandardScaler(with_mean=0, with_std=1).fit_transform(Xt) clf = KNeighborsClassifier(n_neighbors=1) clf.fit(Xs, Ys) ypred = clf.predict(Xt) acc = accuracy_score(y_true=Yt, y_pred=ypred) print('Acc: {:.4f}'.format(acc)) return ypred, acc if __name__ == '__main__': torch.manual_seed(10) # Load data print('Loading data...') data_folder = args.dataset_path domain = {'src': str(args.source), 'tar': str(args.target)} dataloaders = {} data_test = data_load.load_data( data_folder+domain['tar'] + 'images/', BATCH_SIZE['tar'], 'test') data_train = data_load.load_data( data_folder+domain['src'] + 'images/', BATCH_SIZE['src'], 'train', train_val_split=True, train_ratio=.8) dataloaders['train'], dataloaders['val'], dataloaders['test'] = data_train[0], data_train[1], data_test print('Data loaded: Source: {}, Target: {}'.format(args.source, args.target)) # Finetune if args.finetune == 1: print('Begin fintuning...') net = models.Network(base_net=args.model_name, n_class=args.num_class).to(DEVICE) criterion = nn.CrossEntropyLoss() optimizer = get_optimizer(net) if not os.path.exists('save_model/'): os.makedir('save_model/') save_path = 'save_model/best_{}_{}.pth'.format( args.model_name, args.source) model_best, best_acc, acc_hist = finetune( net, dataloaders, optimizer, criterion, save_path, use_lr_schedule=False) print('Finetune completed!') # Extract features from finetuned model if args.extract == 1: model_path = 'save_model/best_{}_{}.pth'.format( args.model_name, args.source) feature_save_path = 'save_model/{}_{}_{}.csv'.format( args.source, args.target, args.model_name) print(feature_save_path) extract_feature( None, dataloaders['test'], feature_save_path, load_from_disk=True, model_path=model_path) print('Deep features are extracted and saved!') # After features are extracted, you may want to classify. # You can try the classify_1nn() function. For instance: # classify_1nn('save_model/amazon_amazon_resnet50.csv', 'save_model/amazon_webcam_resnet50.csv')