def main():
start = time.time()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# load image
img_path = "../tulip_1.jpg"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
json_path = './class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(
json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create model
model = vgg(model_name="vgg16", num_classes=5).to(device)
# load model weights
weights_path = "./vgg16Net.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(
weights_path)
model.load_state_dict(torch.load(weights_path, map_location=device))
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print("预测时间为: {}".format(time.time() - start))
print_res = "class: {} prob: {:.3}".format(
class_indict[str(predict_cla)], predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
def main():
#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
print("using {} device.".format(device))
data_transform = {
"train":
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]),
"val":
transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
}
data_root = os.path.abspath(os.path.join(os.getcwd(),
"../")) # get data root path
image_path = os.path.join(data_root, "data_set",
"flower_data") # flower data set path
assert os.path.exists(image_path), "{} path does not exist.".format(
image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(
image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
#nw = 0
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(
image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(
train_num, val_num))
# test_data_iter = iter(validate_loader)
# test_image, test_label = test_data_iter.next()
# ----加载预训练网络权重
net = vgg()
model_weight_path = "./vgg16-pre.pth" # 预训练权重
assert os.path.exists(model_weight_path), "file {} does not exist.".format(
model_weight_path)
net.load_state_dict(torch.load(model_weight_path, map_location=device))
# ----修改全连接层
'''
in_channel = net.fc.in_features
net.fc = nn.Linear(in_channel, 5)
'''
net.classifier = torch.nn.Sequential(nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(True), nn.Dropout(p=0.5),
nn.Linear(4096, 4096), nn.ReLU(True),
nn.Dropout(p=0.5), nn.Linear(4096, 5))
net.to(device) # 再看这句
# define loss function
loss_function = nn.CrossEntropyLoss()
# construct an optimizer
params = [p for p in net.parameters() if p.requires_grad]
optimizer = optim.Adam(params, lr=0.0001)
# ------开始训练
epochs = 3 # 迁移学习少跑几个循环
best_acc = 0.0
save_path = './vgg16Net.pth' # 保存自己训练的模型
train_steps = len(train_loader)
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(
epoch + 1, epochs, loss)
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')
def main(model_name, feats_name, model_type):
"""Main function."""
compatibility_file = 'data/label/fashion_compatibility_prediction.txt'
data = h5py.File(feats_name, 'r')
data_dict = dict()
for fname, feat in zip(data['filenames'], data['features']):
data_dict[fname] = feat
if model_type == 'inception':
model = inception(512, 512, 2480, batch_first=True, dropout=0.7)
elif model_type == 'vgg':
model = vgg(512, 512, 2480, batch_first=True, dropout=0.7)
elif model_type == 'squeezenet':
model = squeezenet(512, 512, 2480, batch_first=True, dropout=0.7)
else:
print(
"Please, specify a valid model type: inception, vgg or squeezenet, instead of %s"
% model_type)
evaluator = Evaluation(model,
model_type,
model_name,
'data/images',
batch_first=True,
cuda=True)
seqs = [l.replace('\n', '') for l in open(compatibility_file).readlines()]
labels = []
scores = []
tic = time.time()
feats = np.zeros(
(np.sum([len(s.split()[1:]) for s in seqs]) + 2 * len(seqs),
list(data_dict.values())[2].shape[0]))
seq_start_idx = [None] * len(seqs)
seq_start_idx[0] = 0
cum_lens = np.hstack((0, np.cumsum([len(s.split()[1:]) for s in seqs])))
for i, seq in enumerate(seqs):
sys.stdout.write("Concatenating sequences (%d/%d)\r" % (i, len(seqs)))
sys.stdout.flush()
if i != 0:
seq_start_idx[i] = cum_lens[i] + 2 * i
seqimgs = seq.split()[1:]
# Disable complaints about no-member in torch
# pylint: disable=E1101
try:
im_feats = torch.from_numpy(
np.array([data_dict[d] for d in seqimgs]))
except KeyError as e:
im_feats = torch.from_numpy(
np.array([data_dict[bytes(d, 'utf8')] for d in seqimgs]))
# print("Key %s not in the precomputed features." % e)
# import epdb
# epdb.set_trace()
feats[seq_start_idx[i] + 1:seq_start_idx[i] + 1 +
len(im_feats), :] = im_feats
feats = torch.from_numpy(feats.astype(np.float32))
feats = torch.nn.functional.normalize(feats, p=2, dim=1)
for i, seq in enumerate(seqs):
seqtag = seq.split()[0]
seqdata = seq.split()[1:]
compat = evaluator.compatibility(seqdata, data_dict, feats,
seq_start_idx[i])
scores.append(compat)
labels.append(int(seqtag))
sys.stdout.write(
"(%d/%d) SEQ LENGTH = %d - TAG: %s - COMPAT: %.4f - %.2f min left\r"
% (i, len(seqs), len(seqdata), seqtag, compat,
(time.time() - tic) / (i + 1) * (len(seqs) - i) / 60))
sys.stdout.flush()
fpr, tpr, _ = metrics.roc_curve(labels, scores, pos_label=1)
print("\033[0;31m\nModel: %s\033[0m" % model_name)
print("\033[1;30mCompatibility AUC: %f for %d outfits\033[0m" %
(metrics.auc(fpr, tpr), len(labels)))
def get_features(model_name, feats_filename, model_type):
"""Main function for feature extraction."""
if not os.path.exists(feats_filename):
batch_size = 1
if model_type == 'inception':
model = inception(512, 512, 2480, batch_first=True, dropout=0.7)
elif model_type == 'vgg':
model = vgg(512, 512, 2480, batch_first=True, dropout=0.7)
elif model_type == 'squeezenet':
model = squeezenet(512, 512, 2480, batch_first=True, dropout=0.7)
else:
print("Please, specify a valid model type: inception, vgg or squeezenet"\
"instead of %s" % model_type)
return
evaluator = Evaluation(model, model_type, model_name, 'data/images',
batch_first=True, cuda=True)
json_filenames = {'train': 'train_no_dup.json',
'test': 'test_no_dup.json',
'val': 'valid_no_dup.json'}
img_dir, json_dir = 'data/images', 'data/label'
dataloaders = {x: torch.utils.data.DataLoader(
PolyvoreDataset(os.path.join(json_dir, json_filenames[x]), img_dir,
img_transform=None, txt_transform=None),
batch_size=batch_size,
shuffle=False, num_workers=4,
collate_fn=collate_seq)
for x in ['test']}
test_files = json.load(open(os.path.join(json_dir, json_filenames['test'])))
filenames = []
features = torch.Tensor().cuda()
for i, (test_file, batch) in enumerate(zip(test_files, dataloaders['test'])):
if i == 0:
tic = time.time()
sys.stdout.write("%d/%d sets - %.2f secs remaining\r" % (i, len(test_files),
(time.time() - tic)/
(i + 1)*(len(test_files) - i)))
sys.stdout.flush()
set_id = test_file['set_id']
im_idxs = [x['index'] for x in test_file['items']]
im_feats = evaluator.get_img_feats(batch[0]['images'])
for idx in im_idxs:
filenames.append(set_id + '_' + str(idx))
features = torch.cat((features, im_feats.data))
for ignored in batch[0]['ignored']:
filenames.remove(ignored)
if not os.path.exists(os.path.dirname(feats_filename)):
os.makedirs(os.path.dirname(feats_filename))
filenames = [n.encode("ascii", "ignore") for n in filenames]
savefile = h5py.File(feats_filename, 'w')
savefile.create_dataset('filenames', data=filenames)
savefile.create_dataset('features', data=features)
savefile.close()