def main():
"""Execute main function."""
# read parameters for wanted dataset from config file
with open('config.json') as json_file:
config_json = json.load(json_file)
config = config_json[args['dataset']]
if args['mode'] in ['SVM', 'fine_tuning']:
config_source = config_json[args['source_dataset']]
# create directories to save results if they don't exist yet
resultspath = os.path.join(os.path.dirname(os.getcwd()), 'results')
if not os.path.exists(resultspath):
os.makedirs(resultspath)
if not os.path.exists(config['model_savepath']):
os.makedirs(config['model_savepath'])
if not os.path.exists(config['plot_path']):
os.makedirs(config['plot_path'])
# assign batch size
batchsize = int(args['batchsize'])
# set a random seed
seed = 28
if args['mode'] == 'from_scratch':
# set random seed for result reproducability
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
random.seed(seed)
# tf.set_random_seed(seed)
tf.random.set_seed(seed)
# set parameters for training
learning_rate = 5e-6
epochs = 100
# load VGG16 model architecture
model = VGG16(dropout_rate=0.3,
l2_rate=0.0,
batchnorm=True,
activation='relu',
input_shape=(224, 224, 3)).get_model()
model.summary()
# create network instance
network = NeuralNetwork(model, config, batchsize=batchsize, seed=seed)
# compile network
print("compiling network...")
network.compile_network(learning_rate,
optimizer='adam',
loss="binary_crossentropy",
metrics=["accuracy"])
# train network
print("training network...")
history = network.train(epochs)
# save history
print("saving training history...")
network.save_history(history)
# save network
print("saving network...")
network.save_model()
# plot training
print("plotting training progress...")
network.plot_training(history)
# evaluate network on test data
print("evaluating network on test data...")
network.evaluate(mode=args['mode'])
if args['mode'] == 'SVM':
# load the pre-trained source network
print("loading source network...")
source_dataset = args['source_dataset']
modelpath = os.path.join(config_source['model_savepath'],
'{}_model.h5'.format(source_dataset))
source_model = load_model(modelpath)
source_model.summary()
# create network instance
network = NeuralNetwork(source_model,
config,
batchsize=batchsize,
seed=seed)
# set create a bottleneck model at specified layer
network.set_bottleneck_model(outputlayer='flatten_1')
network.model.summary()
# extract features using bottleneck model
bn_features_train, bn_features_test, true_labels_train, true_labels_test = network.extract_bottleneck_features(
)
# scale the data to zero mean, unit variance for PCA
scaler = StandardScaler()
train_features = scaler.fit_transform(bn_features_train)
test_features = scaler.transform(bn_features_test)
# determine time taken to perform PCA and train SVM
start = time()
# fit PCA
print("performing PCA...")
pca = PCA(.95)
pca.fit(train_features)
# apply PCA to features and test data
reduced_train_features = pca.transform(train_features)
reduced_test_features = pca.transform(test_features)
# fit SVM classifier
print("fitting SVM classifier...")
clf = svm.LinearSVC(penalty='l2', loss='squared_hinge',
C=1.0).fit(reduced_train_features,
true_labels_train)
# determine time taken to perform PCA and train SVM
end = time()
training_time = end - start
# make predictions using the trained SVM
preds = clf.decision_function(reduced_test_features)
# calculate AUC and sklearn AUC
print("evaluating results...")
fpr, tpr, thresholds, AUC = AUC_score(preds, true_labels_test)
# calculate accuracy score
acc = accuracy(preds, true_labels_test)
# create a savepath for results and create a sheet to avoid errors
savefile = os.path.join(config['output_path'], 'results_g.xlsx')
# also create excel file if it doesn't exist yet to avoid errors
if not os.path.exists(savefile):
Workbook().save(savefile)
# save in path for target dataset
test_results = pd.Series({
'source_dataset': source_dataset,
'AUC': AUC,
'acc': acc,
'training_time': training_time
})
# read existing results if they exist, else create empty dataframe
try:
df = pd.read_excel(savefile, sheet_name='SVM')
except:
df = pd.DataFrame(
columns=['source_dataset', 'AUC', 'acc', 'training_time'])
df = df.append(test_results, ignore_index=True)
df.set_index('source_dataset', inplace=True)
# save results in excel file
with pd.ExcelWriter(savefile, engine='openpyxl') as writer:
writer.book = load_workbook(savefile)
writer.sheets = dict(
(ws.title, ws) for ws in writer.book.worksheets)
df.index.name = 'source_dataset'
df.to_excel(writer, sheet_name='SVM')
if args['mode'] == 'fine_tuning':
# load the pre-trained source network
print("loading source network...")
source_dataset = args['source_dataset']
modelpath = os.path.join(config_source['model_savepath'],
'{}_model.h5'.format(source_dataset))
source_model = load_model(modelpath)
source_model.summary()
# create network instance
network = NeuralNetwork(source_model,
config,
batchsize=batchsize,
seed=seed)
# create a bottleneck model at specified layer
network.set_bottleneck_model(outputlayer='flatten_1')
network.model.summary()
# freeze all layers in the convolutional base to exclude them from training
for layer in network.model.layers:
layer.trainable = False
# set learning phase to avoid problems with BN layers freezing/training
K.set_learning_phase(1)
# add classification model on top of the bottleneck model
network.set_ft_model()
network.model.summary()
# compile network
print("compiling network...")
network.compile_network(learning_rate=1e-3,
optimizer='sgd',
loss="binary_crossentropy",
metrics=["accuracy"])
# train the model for some epochs
print("training top model...")
epochs = 25
history = network.train(epochs)
histsavepath = os.path.join(config['model_savepath'],
"{}_history_fc".format(args['dataset']))
with open(histsavepath, 'wb') as file:
pickle.dump(history, file)
# evaluate on test data
print("evaluating top model...")
network.evaluate(mode='fc', source_dataset=source_dataset)
# find correct layer index for last conv block and unfreeze last convolutional block
for idx, layer in enumerate(network.model.layers):
if layer.name == 'conv2d_11':
index = idx
for layer in network.model.layers[index:]:
layer.trainable = True
# reset training and validation generators
network.gen_training.reset()
network.gen_validation.reset()
# recompile model
print("compiling network...")
network.compile_network(learning_rate=1e-4,
optimizer='sgd',
loss="binary_crossentropy",
metrics=["accuracy"])
# train model some more
print("training top model and unfrozen layers...")
epochs = 25
history = network.train(epochs)
histsavepath = os.path.join(
config['model_savepath'],
"{}_history_finetuning".format(args['dataset']))
with open(histsavepath, 'wb') as file:
pickle.dump(history, file)
# evaluate results on test data
print("evaluating after fine-tuning top model...")
network.evaluate(mode='fine_tuning', source_dataset=source_dataset)
if args['mode'] == 'evaluate':
# load the source network
print("loading source network...")
modelpath = os.path.join(config['model_savepath'],
'{}_model.h5'.format(args['dataset']))
model = load_model(modelpath)
model.summary()
# create network instance
network = NeuralNetwork(model, config, batchsize=batchsize, seed=seed)
# fit generators on training data
x_train = load_training_data(network.trainingpath)
network.gen_obj_training.fit(x_train, seed=seed)
network.gen_obj_test.fit(x_train, seed=seed)
# initialize generators
network.init_generators(shuffle_training=False,
shuffle_validation=False,
shuffle_test=False,
batchsize=1)
# evaluate network on test data
print("evaluating network on test data...")
network.evaluate(mode=args['mode'])
args = parser.parse_args()
style_name = args.style_image.split("/")[-1].split(".")[0]
os.makedirs(f"images/outputs/{style_name}-training", exist_ok=True)
os.makedirs(f"checkpoints", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Create dataloader for the training data
train_dataset = datasets.ImageFolder(args.dataset_path,
train_transform(args.image_size))
dataloader = DataLoader(train_dataset, batch_size=args.batch_size)
# Defines networks
transformer = TransformerNet().to(device)
vgg = VGG16(requires_grad=False).to(device)
# Load checkpoint model if specified
if args.checkpoint_model:
transformer.load_state_dict(torch.load(args.checkpoint_model))
# Define optimizer and loss
optimizer = Adam(transformer.parameters(), args.lr)
l2_loss = torch.nn.MSELoss().to(device)
# Load style image
style = style_transform(args.style_size)(Image.open(args.style_image))
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
# Extract style features
features_style = vgg(style)
def __init__(self,
generator: Union[Generator, nn.DataParallel],
discriminator: Union[Discriminator, nn.DataParallel],
training_dataset: DataLoader,
validation_dataset: Dataset,
validation_dataset_fid: DataLoader,
vgg16: Union[VGG16, nn.DataParallel] = VGG16(),
generator_optimizer: torch.optim.Optimizer = None,
discriminator_optimizer: torch.optim.Optimizer = None,
generator_loss: nn.Module = LSGANGeneratorLoss(),
discriminator_loss: nn.Module = LSGANDiscriminatorLoss(),
semantic_reconstruction_loss: nn.Module = SemanticReconstructionLoss(),
diversity_loss: nn.Module = DiversityLoss(),
save_data_path: str = 'saved_data') -> None:
'''
Constructor
:param generator: (nn.Module, nn.DataParallel) Generator network
:param discriminator: (nn.Module, nn.DataParallel) Discriminator network
:param training_dataset: (DataLoader) Training dataset
:param vgg16: (nn.Module, nn.DataParallel) VGG16 module
:param generator_optimizer: (torch.optim.Optimizer) Optimizer of the generator network
:param discriminator_optimizer: (torch.optim.Optimizer) Optimizer of the discriminator network
:param generator_loss: (nn.Module) Generator loss function
:param discriminator_loss: (nn.Module) Discriminator loss function
:param semantic_reconstruction_loss: (nn.Module) Semantic reconstruction loss function
:param diversity_loss: (nn.Module) Diversity loss function
'''
# Save parameters
self.generator = generator
self.discriminator = discriminator
self.training_dataset = training_dataset
self.validation_dataset = validation_dataset
self.validation_dataset_fid = validation_dataset_fid
self.vgg16 = vgg16
self.generator_optimizer = generator_optimizer
self.discriminator_optimizer = discriminator_optimizer
self.generator_loss = generator_loss
self.discriminator_loss = discriminator_loss
self.semantic_reconstruction_loss = semantic_reconstruction_loss
self.diversity_loss = diversity_loss
self.latent_dimensions = self.generator.module.latent_dimensions \
if isinstance(self.generator, nn.DataParallel) else self.generator.latent_dimensions
# Init logger
self.logger = Logger()
# Make directories to save logs, plots and models during training
time_and_date = str(datetime.now())
self.path_save_models = os.path.join(save_data_path, 'models_' + time_and_date)
if not os.path.exists(self.path_save_models):
os.makedirs(self.path_save_models)
self.path_save_plots = os.path.join(save_data_path, 'plots_' + time_and_date)
if not os.path.exists(self.path_save_plots):
os.makedirs(self.path_save_plots)
self.path_save_metrics = os.path.join(save_data_path, 'metrics_' + time_and_date)
if not os.path.exists(self.path_save_metrics):
os.makedirs(self.path_save_metrics)
# Make indexes for validation plots
validation_plot_indexes = np.random.choice(range(len(self.validation_dataset_fid.dataset)), 49, replace=False)
# Plot and save validation images used to plot generated images
self.validation_images_to_plot, _, self.validation_masks = image_label_list_of_masks_collate_function(
[self.validation_dataset_fid.dataset[index] for index in validation_plot_indexes])
torchvision.utils.save_image(misc.normalize_0_1_batch(self.validation_images_to_plot),
os.path.join(self.path_save_plots, 'validation_images.png'), nrow=7)
# Plot masks
torchvision.utils.save_image(self.validation_masks[0],
os.path.join(self.path_save_plots, 'validation_masks.png'),
nrow=7)
# Generate latents for validation
self.validation_latents = torch.randn(49, self.latent_dimensions, dtype=torch.float32)
# Log hyperparameter
self.logger.hyperparameter['generator'] = str(self.generator)
self.logger.hyperparameter['discriminator'] = str(self.discriminator)
self.logger.hyperparameter['vgg16'] = str(self.vgg16)
self.logger.hyperparameter['generator_optimizer'] = str(self.generator_optimizer)
self.logger.hyperparameter['discriminator_optimizer'] = str(self.discriminator_optimizer)
self.logger.hyperparameter['generator_loss'] = str(self.generator_loss)
self.logger.hyperparameter['discriminator_loss'] = str(self.discriminator_loss)
self.logger.hyperparameter['diversity_loss'] = str(self.diversity_loss)
self.logger.hyperparameter['discriminator_loss'] = str(self.semantic_reconstruction_loss)
('cycle_ratio', args.cycle_ratio),
('num_classes', args.num_classes)])
return config
config = parse_args()
### call data ###
cifar10 = Cifar10()
n_samples = cifar10.num_examples
n_test_samples = cifar10.num_test_examples
### call models ###
model = VGG16(config['num_classes'])
### make folder ###
mother_folder = config['model_name']
try:
os.mkdir(mother_folder)
except OSError:
pass
### outputs ###
pred, loss = model.Forward()
iter_per_epoch = int(n_samples/config['batch_size'])
mode = None
raise Exception('false mode input')
weightsPath = args.weights_path
Device = torch.device("cuda")
# model select
if modelName == 'AlexNet':
from models.AlexNet import *
model = AlexNet()
elif modelName == "SqueezeNet":
from models.SqueezeNet import *
model = SqueezeNet()
elif modelName == "VGG16":
from models.VGG16 import *
model = VGG16()
elif modelName == "GoogLeNet":
from models.GoogLeNet import *
model = GoogLeNet()
elif modelName == "ResNet18":
from models.ResNet18 import *
model = ResNet18()
elif modelName == "DenseNet121":
from models.DenseNet import *
model = DenseNet121()
elif modelName == "MobileNet":
from models.MobileNet import *
model = MobileNet()
else:
model = None
def main():
n_style = len(cfg.STYLE_LAYERS)
n_content = len(cfg.CONTENT_LAYERS)
# Load and process style and content images
style_img = Image.open(cfg.style_img_path)
content_img = Image.open(cfg.content_img_path)
style_img, content_img = preprocess([style_img, content_img])
style_img, content_img = style_img.to(cfg.device), content_img.to(
cfg.device)
# Image to be optimized
opt_img = Variable(content_img.data.clone(), requires_grad=True)
# Load model and features
model = VGG16().to(cfg.device)
model.load_state_dict(torch.load(cfg.model_path))
style_targets, content_targets = make_targets(style_img, content_img,
model)
# Make losses and optimizer
style_losses = [GramMSELoss().to(cfg.device)] * n_style
content_losses = [nn.MSELoss().to(cfg.device)] * n_content
optimizer = optim.LBFGS([opt_img])
# Regroup all losses / targets / layers into one list
all_layers = cfg.STYLE_LAYERS + cfg.CONTENT_LAYERS
all_weights = cfg.STYLE_WEIGHTS + cfg.CONTENT_WEIGHTS
all_losses = style_losses + content_losses
all_targets = style_targets + content_targets
# Main loop & optimization
n_iter = [0]
start = time.time()
while n_iter[0] < cfg.max_iter:
def closure():
optimizer.zero_grad()
output = model(opt_img, cfg.STYLE_LAYERS + cfg.CONTENT_LAYERS)
losses_values = [
all_weights[i] * all_losses[i](output[i], all_targets[i])
for i in range(n_style + n_content)
]
total_loss = sum(losses_values)
total_loss.backward()
n_iter[0] += 1
if n_iter[0] % cfg.show_every == 0:
print("[{}/{}] loss: {:.4f} elapsed time: {:.2f}s".format(
n_iter[0],
cfg.max_iter,
total_loss.item(),
time.time() - start,
))
return total_loss
optimizer.step(closure)
# Post process and save the output image
output_image = opt_img.data[0].cpu().squeeze()
output_image = postprocess(output_image)
plt.figure()
plt.imshow(output_image)
plt.show()
output_image.save(cfg.output_img)
def train(i, train_acc, train_loss):
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
#print(DATASET_ROOT)
train_set = IMAGE_Dataset(Path(DATASET_ROOT1), data_transform)
data_loader = DataLoader(dataset=train_set,
batch_size=32,
shuffle=True,
num_workers=1)
#print(train_set.num_classes)
if i == 1:
model = VGG16(num_classes=train_set.num_classes)
vgg19 = models.vgg19_bn(pretrained=True) #載入vgg16 model
pretrained_dict = vgg19.state_dict() #vgg16預訓練的參數
model_dict = model.state_dict() #自訂VGG16的參數
#vgg16 = nn.Sequential(*list(model.children())[:-1])
# 將pretrained_dict裏不屬於model_dict的鍵剔除掉
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict and (
model_dict[k].shape == pretrained_dict[k].shape)
}
model_dict.update(pretrained_dict) # 更新現有的model_dict
#print(pretrained_dict)
#print(model_dict)
model.load_state_dict(model_dict) # 加載我們真正需要的state_dict
if i != 1:
model = torch.load(PATH_TO_WEIGHTS)
'''if i==1:
model = models.resnet101(pretrained=True)
fc_features=model.fc.in_features
model.fc=nn.Linear(fc_features,196)
if i!=1:
model=torch.load(PATH_TO_WEIGHTS)'''
model = model.cuda(CUDA_DEVICES)
model.train() #train
best_model_params = copy.deepcopy(model.state_dict()) #複製參數
best_acc = 0.0
num_epochs = 10
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model.parameters(),
lr=0.001,
momentum=0.9)
for epoch in range(num_epochs):
print(f'Epoch: {epoch + 1}/{num_epochs}')
print('-' * len(f'Epoch: {epoch + 1}/{num_epochs}'))
training_loss = 0.0
training_corrects = 0
for i, (inputs, labels) in enumerate(data_loader):
inputs = Variable(inputs.cuda(CUDA_DEVICES))
labels = Variable(labels.cuda(CUDA_DEVICES))
optimizer.zero_grad()
outputs = model(inputs)
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
training_loss += loss.item() * inputs.size(0)
#revise loss.data[0]-->loss.item()
training_corrects += torch.sum(preds == labels.data)
#print(f'training_corrects: {training_corrects}')
#if(not(i%10)):
# print(f'iteration done :{i}')
training_loss = training_loss / len(train_set) #train loss
training_acc = training_corrects.double() / len(train_set) #tarin acc
#print(training_acc.type())
#print(f'training_corrects: {training_corrects}\tlen(train_set):{len(train_set)}\n')
print(
f'Training loss: {training_loss:.4f}\taccuracy: {training_acc:.4f}\n'
)
train_acc.append(training_acc) #save each 10 epochs accuracy
train_loss.append(training_loss)
if training_acc > best_acc:
best_acc = training_acc
best_model_params = copy.deepcopy(model.state_dict())
model.load_state_dict(
best_model_params) #model load new best parms #model載入參數
torch.save(model, f'model-best_train_acc.pth') #save model #存整個model
return (train_acc, train_loss)
rgt = gt_rgb rgt_sz = [prm.batch_size,prm.crop_size,prm.crop_size,3] rgt_b0 = tf.Variable(tf.zeros(rgt_sz,dtype=tf.float32)) rgt_b1 = tf.Variable(tf.zeros(rgt_sz,dtype=tf.float32)) tldr_fetchOp = [rinp_b0.assign(rinp).op, rgt_b0.assign(rgt).op] vldr_fetchOp = [rinp_b1.assign(rinp).op, rgt_b1.assign(rgt).op] tldr_swapOp = [rinp_b1.assign(rinp_b0).op, rgt_b1.assign(rgt_b0).op] # Init RefineNet R = RefineNet(rinp_b1,bn='train',outp_act=False) rpred = (R.pred+1.)*127.5 # Init perceptual network pinp = tf.concat((rgt_b1,rpred),axis=0) P = VGG16(pinp,stop_layer='conv3_3') ppred = P.pred # Init discriminator network dgt0 = tf.constant(0,shape=[prm.batch_size],dtype=tf.int64) dgt1 = tf.constant(1,shape=[prm.batch_size],dtype=tf.int64) dgt = tf.concat((dgt0,dgt1),axis=0) layers = ['conv1_1','conv2_2','conv3_3'] dinp_fake = [ppred[layer][prm.batch_size:] for layer in layers] dinp_real = [ppred[layer][:prm.batch_size] for layer in layers] dinp_fake[0] = tf.concat((rinp_b1,rpred,dinp_fake[0]),axis=3) dinp_real[0] = tf.concat((rinp_b1,rgt_b1,dinp_real[0]),axis=3) D = Discriminator() dpred_fake = D.pred(dinp_fake)
STD = np.array([0.229, 0.224, 0.225])
mask_transform = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.Resize(256),
transforms.ToTensor(),
])
device = torch.device(DEVICE if torch.cuda.is_available() else "cpu")
d_net = Discriminator()
refine_net = Net()
if MODE != "ablation":
vgg = VGG16(requires_grad=False)
vgg.to(device)
if torch.cuda.device_count() > 1 and MULTI_GPU:
print("Using {} GPUs...".format(torch.cuda.device_count()))
d_net = nn.DataParallel(d_net)
refine_net = nn.DataParallel(refine_net)
else:
print("GPU ID: {}".format(device))
d_net = d_net.to(device)
refine_net = refine_net.to(device)
if MODE == "train":
refine_net.apply(weights_init)
def main():
global args, best_prec1
args = parser.parse_args()
print(args)
# create model
model = VGG16('pre_trained_models/vgg_places_365.pt', return_output=True)
model = torch.nn.DataParallel(model).cuda()
print(model)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.Scale(256),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Scale(256),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size // 2,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# define loss function (criterion) and pptimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.Adam(model.parameters(), args.lr)
if args.evaluate:
validate(val_loader, model, criterion)
return
# evaluate on validation set
validate(val_loader, model, criterion)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best, "VGG_16")
def train():
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
#print(DATASET_ROOT)
train_set = IMAGE_Dataset(Path(DATASET_ROOT), data_transform)
data_loader = DataLoader(dataset=train_set,
batch_size=32,
shuffle=True,
num_workers=1)
#print(train_set.num_classes)
model = VGG16(num_classes=train_set.num_classes)
model = model.cuda(CUDA_DEVICES)
model.train()
best_model_params = copy.deepcopy(model.state_dict())
best_acc = 0.0
count = 100.0
num_epochs = 35
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model.parameters(),
lr=0.01,
momentum=0.9)
for epoch in range(num_epochs):
print(f'Epoch: {epoch + 1}/{num_epochs}')
print('-' * len(f'Epoch: {epoch + 1}/{num_epochs}'))
adjust_learning_rate(optimizer, num_epochs)
training_loss = 0.0
training_corrects = 0
for i, (inputs, labels) in enumerate(data_loader):
inputs = Variable(inputs.cuda(CUDA_DEVICES))
labels = Variable(labels.cuda(CUDA_DEVICES))
optimizer.zero_grad()
outputs = model(inputs)
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
training_loss += loss.item() * inputs.size(0)
#revise loss.data[0]-->loss.item()
training_corrects += torch.sum(preds == labels.data)
#print(f'training_corrects: {training_corrects}')
training_loss = training_loss / len(train_set)
training_acc = training_corrects.double() / len(train_set)
result_train_loss.append(str(training_loss)) #train_data
result_train_acc.append(str(training_acc))
if (epoch + 1) % 10 == 0:
#torch.save(model, f'model.pth')
#tmp1,tmp2=test.test(f'model.pth')
#print(str(tmp1)+'\t'+str(tmp2))
torch.save(model, f'model-{count:.02f}-train.pth')
#result_test_loss.append(str(tmp1))
#result_test_acc.append(str(tmp2))
count = count + 1
#print(training_acc.type())
#print(f'training_corrects: {training_corrects}\tlen(train_set):{len(train_set)}\n')
print(
f'Training loss: {training_loss:.4f}\taccuracy: {training_acc:.4f}\n'
)
if training_acc > best_acc:
best_acc = training_acc
best_model_params = copy.deepcopy(model.state_dict())
model.load_state_dict(best_model_params)
torch.save(model, f'model-{best_acc:.02f}-best_train_acc.pth')
def main(args):
num_classes = 8
size = [224, 224] # size of images
in_ = eddl.Input([3, size[0], size[1]])
out = VGG16(in_, num_classes)
net = eddl.Model([in_], [out])
eddl.build(net, eddl.sgd(0.001, 0.9), ["soft_cross_entropy"],
["categorical_accuracy"],
eddl.CS_GPU([1]) if args.gpu else eddl.CS_CPU())
eddl.summary(net)
eddl.setlogfile(net, "skin_lesion_classification_inference")
training_augs = ecvl.SequentialAugmentationContainer([
ecvl.AugResizeDim(size),
])
test_augs = ecvl.SequentialAugmentationContainer([
ecvl.AugResizeDim(size),
])
dataset_augs = ecvl.DatasetAugmentations([training_augs, None, test_augs])
print("Reading dataset")
d = ecvl.DLDataset(args.in_ds, args.batch_size, dataset_augs)
if args.out_dir:
for c in d.classes_:
os.makedirs(os.path.join(args.out_dir, c), exist_ok=True)
x = Tensor([args.batch_size, d.n_channels_, size[0], size[1]])
y = Tensor([args.batch_size, len(d.classes_)])
d.SetSplit(ecvl.SplitType.test)
num_samples = len(d.GetSplit())
num_batches = num_samples // args.batch_size
metric = eddl.getMetric("categorical_accuracy")
total_metric = []
if not os.path.exists(args.ckpts):
raise RuntimeError('Checkpoint "{}" not found'.format(args.ckpts))
eddl.load(net, args.ckpts, "bin")
print("Testing")
for b in range(num_batches):
n = 0
print("Batch {:d}/{:d}".format(b + 1, num_batches))
d.LoadBatch(x, y)
x.div_(255.0)
eddl.forward(net, [x])
output = eddl.getOutput(out)
sum_ = 0.0
for j in range(args.batch_size):
result = output.select([str(j)])
target = y.select([str(j)])
ca = metric.value(target, result)
total_metric.append(ca)
sum_ += ca
if args.out_dir:
result_a = np.array(result, copy=False)
target_a = np.array(target, copy=False)
classe = np.argmax(result_a).item()
gt_class = np.argmax(target_a).item()
single_image = x.select([str(j)])
img_t = ecvl.TensorToView(single_image)
img_t.colortype_ = ecvl.ColorType.BGR
single_image.mult_(255.)
filename = d.samples_[d.GetSplit()[n]].location_[0]
head, tail = os.path.splitext(os.path.basename(filename))
bname = "%s_gt_class_%s.png" % (head, gt_class)
cur_path = os.path.join(args.out_dir, d.classes_[classe],
bname)
ecvl.ImWrite(cur_path, img_t)
n += 1
print("categorical_accuracy:", sum_ / args.batch_size)
total_avg = sum(total_metric) / len(total_metric)
print("Total categorical accuracy:", total_avg)
def VGGFace(include_top=True,
model='vgg16',
weights='vggface',
input_tensor=None,
input_shape=None,
pooling=None,
classes=None):
"""Instantiates the VGGFace architectures.
Optionally loads weights pre-trained
on VGGFace datasets. Note that when using TensorFlow,
for best performance you should set
`image_data_format="channels_last"` in your Keras config
at ~/.keras/keras.json.
The model and the weights are compatible with both
TensorFlow and Theano. The data format
convention used by the model is the one
specified in your Keras config file.
# Arguments
include_top: whether to include the 3 fully-connected
layers at the top of the network.
weights: one of `None` (random initialization)
or "vggface" (pre-training on VGGFACE datasets).
input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
to use as image input for the model.
model: selects the one of the available architectures
vgg16, resnet50 or senet50 default is vgg16.
input_shape: optional shape tuple, only to be specified
if `include_top` is False (otherwise the input shape
has to be `(224, 224, 3)` (with `channels_last` data format)
or `(3, 224, 244)` (with `channels_first` data format).
It should have exactly 3 inputs channels,
and width and height should be no smaller than 48.
E.g. `(200, 200, 3)` would be one valid value.
pooling: Optional pooling mode for feature extraction
when `include_top` is `False`.
- `None` means that the output of the model will be
the 4D tensor output of the
last convolutional layer.
- `avg` means that global average pooling
will be applied to the output of the
last convolutional layer, and thus
the output of the model will be a 2D tensor.
- `max` means that global max pooling will
be applied.
classes: optional number of classes to classify images
into, only to be specified if `include_top` is True, and
if no `weights` argument is specified.
# Returns
A Keras model instance.
# Raises
ValueError: in case of invalid argument for `weights`,
or invalid input shape.
"""
if weights not in {'vggface', None}:
raise ValueError('The `weights` argument should be either '
'`None` (random initialization) or `vggface`'
'(pre-training on VGGFace Datasets).')
if model == 'vgg16':
if classes is None:
classes = 2622
if weights == 'vggface' and include_top and classes != 2622:
raise ValueError(
'If using `weights` as vggface original with `include_top`'
' as true, `classes` should be 2622')
return VGG16(include_top=include_top,
input_tensor=input_tensor,
input_shape=input_shape,
pooling=pooling,
weights=weights,
classes=classes)
if model == 'resnet50':
if classes is None:
classes = 8631
if weights == 'vggface' and include_top and classes != 8631:
raise ValueError(
'If using `weights` as vggface original with `include_top`'
' as true, `classes` should be 8631')
return RESNET50(include_top=include_top,
input_tensor=input_tensor,
input_shape=input_shape,
pooling=pooling,
weights=weights,
classes=classes)
if model == 'senet50':
if classes is None:
classes = 8631
if weights == 'vggface' and include_top and classes != 8631:
raise ValueError(
'If using `weights` as vggface original with `include_top`'
' as true, `classes` should be 8631')
return SENET50(include_top=include_top,
input_tensor=input_tensor,
input_shape=input_shape,
pooling=pooling,
weights=weights,
classes=classes)
def main(_):
print("FLAGS values:", tf.app.flags.FLAGS.flag_values_dict())
tf.logging.set_verbosity(tf.logging.INFO)
config = tf.ConfigProto()
if FLAGS.model == 'denoiser':
config.gpu_options.per_process_gpu_memory_fraction = 0.2
else:
config.gpu_options.allow_growth = True
if FLAGS.model == 'denoiser':
from models import Denoiser
model = Denoiser()
with tf.Graph().as_default():
# Prepare graph
with tf.Session(config=config).as_default():
if FLAGS.model == 'inception-v3':
from models import InceptionV3
model = InceptionV3()
elif FLAGS.model == 'vgg-16':
from models import VGG16
model = VGG16()
elif FLAGS.model == 'resnet-50':
from models import ResNet50
model = ResNet50()
elif FLAGS.model == 'jpeg':
from models import JPEG
model = JPEG()
elif FLAGS.model == 'random':
from models import Random
model = Random()
if FLAGS.method != 'uniform':
from surrogate_model import ResNet152
model_s = ResNet152(source_image_size=model.image_size,
use_smoothed_grad=model.use_smoothed_grad)
image_size = model.image_size
# ---Setting hyperparameters---
if FLAGS.norm == 'l2':
epsilon = 1e-3
eps = np.sqrt(epsilon * image_size * image_size * 3)
learning_rate = 2.0 / np.sqrt(image_size * image_size * 3)
else:
epsilon = 0.05
eps = epsilon
learning_rate = 0.005
if model.use_larger_step_size:
ini_sigma = 1e-3
else:
ini_sigma = 1e-4
# -----------------------------
if not model.predictions_is_correct:
l = open(os.path.join(FLAGS.input_dir, 'labels')).readlines()
gts = {}
for i in l:
i = i.strip().split(' ')
gts[i[0]] = int(i[1]) + (model.num_classes - 1000)
success = 0
queries = []
correct = 0
names_images = load_images(FLAGS.input_dir, model.image_size)
for filename, image in names_images:
output_logging = open(os.path.join(FLAGS.output_dir, 'logging'), 'a')
sigma = ini_sigma
np.random.seed(0)
tf.set_random_seed(0)
adv_image = image.copy()
label = model.get_pred(image)
l = model.get_loss(image, label)
print(filename, 'original prediction:', label, 'loss:', l)
if not model.predictions_is_correct:
correct += (label[0] == gts[filename])
if label[0] != gts[filename]:
output_logging.write(filename + ' original misclassified.\n')
output_logging.close()
continue
lr = learning_rate
last_loss = []
total_q = 0
ite = 0
while total_q <= FLAGS.max_queries:
total_q += 1
if FLAGS.show_true and hasattr(model, 'get_grad'):
true = np.squeeze(model.get_grad(adv_image, label))
print("Grad norm", np.sqrt(np.sum(true * true)))
if ite % 2 == 0 and sigma != ini_sigma:
print(
"sigma has been increased before; checking if sigma could be set back to ini_sigma"
)
rand = np.random.normal(size=adv_image.shape)
rand = rand / np.maximum(1e-12,
np.sqrt(np.mean(np.square(rand))))
rand_loss = model.get_loss(adv_image + ini_sigma * rand, label)
total_q += 1
rand = np.random.normal(size=adv_image.shape)
rand = rand / np.maximum(1e-12,
np.sqrt(np.mean(np.square(rand))))
rand_loss2 = model.get_loss(adv_image + ini_sigma * rand,
label)
total_q += 1
if (rand_loss - l)[0] != 0 and (rand_loss2 - l)[0] != 0:
print("set sigma back to ini_sigma")
sigma = ini_sigma
if FLAGS.method != 'uniform':
if model.num_classes < model_s.num_classes:
s_label = label + 1
elif model.num_classes > model_s.num_classes:
s_label = label - 1
else:
s_label = label
prior = np.squeeze(model_s.get_grad(adv_image, s_label))
if FLAGS.show_true and hasattr(model, 'get_grad'):
alpha = np.sum(true * prior) / np.maximum(
1e-12,
np.sqrt(np.sum(true * true) * np.sum(prior * prior)))
print("alpha =", alpha)
prior = prior / np.maximum(1e-12,
np.sqrt(np.mean(np.square(prior))))
if FLAGS.method in ['biased', 'average']:
start_iter = 3
if ite % 10 == 0 or ite == start_iter:
# Estimate norm of true gradient periodically when ite == 0/10/20...;
# since gradient norm may change fast in the early iterations, we also
# estimate the gradient norm when ite == 3.
s = 10
pert = np.random.normal(size=(s, ) + adv_image.shape)
for i in range(s):
pert[i] = pert[i] / np.maximum(
1e-12, np.sqrt(np.mean(np.square(pert[i]))))
eval_points = adv_image + sigma * pert
losses = model.get_loss(eval_points, np.repeat(label, s))
total_q += s
norm_square = np.average(((losses - l) / sigma)**2)
while True:
prior_loss = model.get_loss(adv_image + sigma * prior,
label)
total_q += 1
diff_prior = (prior_loss - l)[0]
if diff_prior == 0:
# Avoid the numerical issue in finite difference
sigma *= 2
print("multiply sigma by 2")
else:
break
est_alpha = diff_prior / sigma / np.maximum(
np.sqrt(np.sum(np.square(prior)) * norm_square), 1e-12)
print("Estimated alpha =", est_alpha)
alpha = est_alpha
if alpha < 0:
prior = -prior
alpha = -alpha
q = FLAGS.samples_per_draw
n = image_size * image_size * 3
d = 50 * 50 * 3
gamma = 3.5
A_square = d / n * gamma
return_prior = False
if FLAGS.method == 'average':
if FLAGS.dataprior:
alpha_nes = np.sqrt(A_square * q / (d + q - 1))
else:
alpha_nes = np.sqrt(q / (n + q - 1))
if alpha >= 1.414 * alpha_nes:
return_prior = True
elif FLAGS.method == 'biased':
if FLAGS.dataprior:
best_lambda = A_square * (
A_square - alpha**2 *
(d + 2 * q - 2)) / (A_square**2 + alpha**4 * d**2 -
2 * A_square * alpha**2 *
(q + d * q - 1))
else:
best_lambda = (1 - alpha**2) * (
1 - alpha**2 *
(n + 2 * q - 2)) / (alpha**4 * n * (n + 2 * q - 2) -
2 * alpha**2 * n * q + 1)
print('best_lambda = ', best_lambda)
if best_lambda < 1 and best_lambda > 0:
lmda = best_lambda
else:
if alpha**2 * (n + 2 * q - 2) < 1:
lmda = 0
else:
lmda = 1
if np.abs(alpha) >= 1:
lmda = 1
print('lambda = ', lmda)
if lmda == 1:
return_prior = True
elif FLAGS.method == 'fixed_biased':
lmda = FLAGS.fixed_const
if not return_prior:
if FLAGS.dataprior:
pert = np.random.normal(size=(q, 50, 50, 3))
pert = np.array([
cv2.resize(pert[i],
adv_image.shape[:2],
interpolation=cv2.INTER_NEAREST)
for i in range(q)
])
else:
pert = np.random.normal(size=(q, ) + adv_image.shape)
for i in range(q):
if FLAGS.method in ['biased', 'fixed_biased']:
pert[i] = pert[i] - np.sum(
pert[i] * prior) * prior / np.maximum(
1e-12, np.sum(prior * prior))
pert[i] = pert[i] / np.maximum(
1e-12, np.sqrt(np.mean(np.square(pert[i]))))
pert[i] = np.sqrt(1 - lmda) * pert[i] + np.sqrt(
lmda) * prior
else:
pert[i] = pert[i] / np.maximum(
1e-12, np.sqrt(np.mean(np.square(pert[i]))))
while True:
eval_points = adv_image + sigma * pert
losses = model.get_loss(eval_points, np.repeat(label, q))
total_q += q
grad = (losses - l).reshape(-1, 1, 1, 1) * pert
grad = np.mean(grad, axis=0)
norm_grad = np.sqrt(np.mean(np.square(grad)))
if norm_grad == 0:
# Avoid the numerical issue in finite difference
sigma *= 5
print("estimated grad == 0, multiply sigma by 5")
else:
break
grad = grad / np.maximum(1e-12,
np.sqrt(np.mean(np.square(grad))))
if FLAGS.method == 'average':
while True:
diff_prior = (
model.get_loss(adv_image + sigma * prior, label) -
l)[0]
total_q += 1
diff_nes = (
model.get_loss(adv_image + sigma * grad, label) -
l)[0]
total_q += 1
diff_prior = max(0, diff_prior)
if diff_prior == 0 and diff_nes == 0:
sigma *= 2
print("multiply sigma by 2")
else:
break
final = prior * diff_prior + grad * diff_nes
final = final / np.maximum(
1e-12, np.sqrt(np.mean(np.square(final))))
print("diff_prior = {}, diff_nes = {}".format(
diff_prior, diff_nes))
elif FLAGS.method == 'fixed_average':
diff_prior = (
model.get_loss(adv_image + sigma * prior, label) -
l)[0]
total_q += 1
if diff_prior < 0:
prior = -prior
final = FLAGS.fixed_const * prior + (
1 - FLAGS.fixed_const) * grad
final = final / np.maximum(
1e-12, np.sqrt(np.mean(np.square(final))))
else:
final = grad
def print_loss(model, direction):
length = [1e-4, 1e-3]
les = []
for ss in length:
les.append((
model.get_loss(adv_image + ss * direction, label) -
l)[0])
print("losses", les)
if FLAGS.show_loss:
if FLAGS.method in ['average', 'fixed_average']:
lprior = model.get_loss(adv_image + lr * prior,
label) - l
print_loss(model, prior)
lgrad = model.get_loss(adv_image + lr * grad,
label) - l
print_loss(model, grad)
lfinal = model.get_loss(adv_image + lr * final,
label) - l
print_loss(model, final)
print(lprior, lgrad, lfinal)
elif FLAGS.method in ['biased', 'fixed_biased']:
lprior = model.get_loss(adv_image + lr * prior,
label) - l
print_loss(model, prior)
lgrad = model.get_loss(adv_image + lr * grad,
label) - l
print_loss(model, grad)
print(lprior, lgrad)
else:
final = prior
if FLAGS.show_true and hasattr(model, 'get_grad'):
if FLAGS.method in ['average', 'fixed_average'
] and not return_prior:
print(
"NES angle =",
np.sum(true * grad) / np.maximum(
1e-12,
np.sqrt(
np.sum(true * true) * np.sum(grad * grad))))
print(
"angle =",
np.sum(true * final) / np.maximum(
1e-12,
np.sqrt(np.sum(true * true) * np.sum(final * final))))
if FLAGS.norm == 'l2':
adv_image = adv_image + lr * final / np.maximum(
1e-12, np.sqrt(np.mean(np.square(final))))
norm = max(1e-12, np.linalg.norm(adv_image - image))
factor = min(1, eps / norm)
adv_image = image + (adv_image - image) * factor
else:
adv_image = adv_image + lr * np.sign(final)
adv_image = np.clip(adv_image, image - eps, image + eps)
adv_image = np.clip(adv_image, 0, 1)
adv_label = model.get_pred(adv_image)
l = model.get_loss(adv_image, label)
print('queries:', total_q, 'loss:', l, 'learning rate:', lr,
'sigma:', sigma, 'prediction:', adv_label, 'distortion:',
np.max(np.abs(adv_image - image)),
np.linalg.norm(adv_image - image))
ite += 1
if adv_label != label:
print('Stop at queries:', total_q)
success += 1
queries.append(total_q)
imsave(os.path.join(FLAGS.output_dir, filename), adv_image)
output_logging.write(filename + ' succeed; queries: ' +
str(total_q) + '\n')
break
else:
imsave(os.path.join(FLAGS.output_dir, filename), adv_image)
output_logging.write(filename + ' fail.\n')
output_logging.close()
if model.predictions_is_correct:
total = FLAGS.number_images
else:
total = correct
print('Success rate:', success / total, 'Queries', queries)
output_logging = open(os.path.join(FLAGS.output_dir, 'logging'), 'a')
output_logging.write('Success rate: ' + str(success / total) +
', Queries on success: ' + str(np.mean(queries)))
output_logging.close()
def train():
parser = argparse.ArgumentParser(description='parser for style transfer')
parser.add_argument('--dataset_path',
type=str,
default=r'C:\Users\Dewey\data\celeba',
help='path to training dataset')
parser.add_argument('--style_image',
type=str,
default='mosaic.jpg',
help='path to style img')
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--batch_size', type=int, default=4)
parser.add_argument('--image_size',
type=int,
default=256,
help='training image size')
parser.add_argument('--style_img_size',
type=int,
default=256,
help='style image size')
parser.add_argument("--lambda_content",
type=float,
default=1e5,
help="Weight for content loss")
parser.add_argument("--lambda_style",
type=float,
default=1e10,
help="Weight for style loss")
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument("--checkpoint_model",
type=str,
help="Optional path to checkpoint model")
parser.add_argument("--checkpoint_interval",
type=int,
default=2000,
help="Batches between saving model")
parser.add_argument("--sample_interval",
type=int,
default=1000,
help="Batches between saving image samples")
parser.add_argument('--sample_format',
type=str,
default='jpg',
help='sample image format')
args = parser.parse_args()
style_name = args.style_image.split('/')[-1].split('.')[0]
os.makedirs(f'images/outputs/{style_name}-training',
exist_ok=True) # f-string格式化字符串
os.makedirs('checkpoints', exist_ok=True)
def save_sample(batch):
transformer.eval()
with torch.no_grad():
output = transformer(image_samples.to(device))
img_grid = denormalize(
torch.cat((image_samples.cpu(), output.cpu()), 2))
save_image(img_grid,
f"images/outputs/{style_name}-training/{batch}.jpg",
nrow=4)
transformer.train()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = datasets.ImageFolder(args.dataset_path,
train_transform(args.image_size))
dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
transformer = TransformerNet().to(device)
vgg = VGG16(requires_grad=False).to(device)
if args.checkpoint_model:
transformer.load_state_dict(torch.load(args.checkpoint_model))
optimizer = Adam(transformer.parameters(), lr=args.lr)
l2_loss = nn.MSELoss().to(device)
# load style image
style = style_transform(args.style_img_size)(Image.open(args.style_image))
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
# style_image features
style_features = vgg(style)
gram_style = [gram(x) for x in style_features]
# visualization the image
image_samples = []
for path in random.sample(
glob.glob(f'{args.dataset_path}/*/*.{args.sample_format}'), 8):
image_samples += [
style_transform(
(args.image_size, args.image_size))(Image.open(path))
]
image_samples = torch.stack(image_samples)
c_loss = 0
s_loss = 0
t_loss = 0
for epoch in range(args.epochs):
for i, (img, _) in enumerate(dataloader):
optimizer.zero_grad()
image_original = img.to(device)
image_transformed = transformer(image_original)
origin_features = vgg(image_original)
transformed_features = vgg(image_transformed)
content_loss = args.lambda_content * l2_loss(
transformed_features.relu_2_2, origin_features.relu_2_2)
style_loss = 0
for ii, jj in zip(transformed_features, gram_style):
gram_t_features = gram(ii)
style_loss += l2_loss(gram_t_features,
jj[:img.size(), :, :]) # buyiyang
style_loss *= args.lambda_style
loss = content_loss + style_loss
loss.backward()
optimizer.step()
c_loss += content_loss.item()
s_loss += style_loss.item()
t_loss += loss.item()
print(
'[Epoch %d/%d] [Batch %d/%d] [Content: %.2f (%.2f) Style: %.2f (%.2f) Total: %.2f (%.2f)]'
% (
epoch + 1,
args.epochs,
i,
len(train_dataset),
content_loss.item(),
np.mean(c_loss),
style_loss.item(),
np.mean(s_loss),
loss.item(),
np.mean(t_loss),
))
batches_done = epoch * len(dataloader) + i + 1
if batches_done % args.sample_interval == 0:
save_sample(batches_done)
if args.checkpoint_interval > 0 and batches_done % args.checkpoint_interval == 0:
style_name = os.path.basename(args.style_image).split(".")[0]
torch.save(transformer.state_dict(),
f"checkpoints/{style_name}_{batches_done}.pth")
condition_features = np.concatenate(condition_features)
return condition_features
if __name__ == '__main__':
parser = ArgumentParser(description='PCA encoder using BOLD5000 study data')
parser.add_argument('--bold5000_folder', required=True, type=str, help='folder containing the stimuli images')
parser.add_argument('--feature_extractor', default='alexnet', type=str, help='feature extraction model')
parser.add_argument('--feature_name', default='conv_3', type=str, help='feature extraction layer')
parser.add_argument('--n_pcs', default=400, type=int, help='number of pcs to fit')
args = parser.parse_args()
if args.feature_extractor == 'alexnet':
feat_extractor = AlexNet(args.feature_name)
elif args.feature_extractor == 'vgg16':
feat_extractor = VGG16(args.feature_name)
else:
raise ValueError('unimplemented feature extractor: {}'.format(args.feature_extractor))
if torch.cuda.is_available():
feat_extractor.cuda()
features = condition_features(os.path.join(args.bold5000_folder, 'stimuli'), feat_extractor)
pca = PCA(n_components=args.n_pcs)
pca.fit(features)
print('\nMean Explained Variance: {:.4f}'.format(np.mean(pca.explained_variance_ratio_.mean())))
encoder = PCAEncoder(feat_extractor, pcs=pca.components_, mean=pca.mean_)
encoder.eval()
run_name = utils.get_run_name('bold5000', args.feature_extractor, args.feature_name, ['PCA'])
torch.save(encoder, os.path.join('saved_models', run_name + '.pth'))