def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
# load image
img_path = "../tulip.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 = create_model(num_classes=5, has_logits=False).to(device)
# load model weights
model_weight_path = "./weights/model-9.pth"
model.load_state_dict(torch.load(model_weight_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_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():
num_classes = 5
im_height = im_width = 224
# load image
img_path = "../tulip.jpg"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
img = Image.open(img_path)
# resize image
img = img.resize((im_width, im_height))
plt.imshow(img)
# read image
img = np.array(img).astype(np.float32)
# preprocess
img = (img / 255. - 0.5) / 0.5
# Add the image to a batch where it's the only member.
img = (np.expand_dims(img, 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 = create_model(num_classes=num_classes, has_logits=False)
model.build([1, 224, 224, 3])
weights_path = './save_weights/model.ckpt'
assert len(glob.glob(weights_path+"*")), "cannot find {}".format(weights_path)
model.load_weights(weights_path)
result = np.squeeze(model.predict(img, batch_size=1))
result = tf.keras.layers.Softmax()(result)
predict_class = np.argmax(result)
print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_class)],
result[predict_class])
plt.title(print_res)
print(print_res)
plt.show()
def main(args):
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
if os.path.exists("./weights") is False:
os.makedirs("./weights")
tb_writer = SummaryWriter()
train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(
args.data_path)
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(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
}
# 实例化训练数据集
train_dataset = MyDataSet(images_path=train_images_path,
images_class=train_images_label,
transform=data_transform["train"])
# 实例化验证数据集
val_dataset = MyDataSet(images_path=val_images_path,
images_class=val_images_label,
transform=data_transform["val"])
batch_size = args.batch_size
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0,
8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=nw,
collate_fn=train_dataset.collate_fn)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
num_workers=nw,
collate_fn=val_dataset.collate_fn)
model = create_model(num_classes=5, has_logits=False).to(device)
if args.weights != "":
assert os.path.exists(
args.weights), "weights file: '{}' not exist.".format(args.weights)
weights_dict = torch.load(args.weights, map_location=device)
# 删除不需要的权重
del_keys = ['head.weight', 'head.bias'] if model.has_logits \
else ['pre_logits.fc.weight', 'pre_logits.fc.bias', 'head.weight', 'head.bias']
for k in del_keys:
del weights_dict[k]
print(model.load_state_dict(weights_dict, strict=False))
if args.freeze_layers:
for name, para in model.named_parameters():
# 除head, pre_logits外，其他权重全部冻结
if "head" not in name and "pre_logits" not in name:
para.requires_grad_(False)
else:
print("training {}".format(name))
pg = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=5E-5)
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (
1 - args.lrf) + args.lrf # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
for epoch in range(args.epochs):
# train
train_loss, train_acc = train_one_epoch(model=model,
optimizer=optimizer,
data_loader=train_loader,
device=device,
epoch=epoch)
scheduler.step()
# validate
val_loss, val_acc = evaluate(model=model,
data_loader=val_loader,
device=device,
epoch=epoch)
tags = [
"train_loss", "train_acc", "val_loss", "val_acc", "learning_rate"
]
tb_writer.add_scalar(tags[0], train_loss, epoch)
tb_writer.add_scalar(tags[1], train_acc, epoch)
tb_writer.add_scalar(tags[2], val_loss, epoch)
tb_writer.add_scalar(tags[3], val_acc, epoch)
tb_writer.add_scalar(tags[4], optimizer.param_groups[0]["lr"], epoch)
torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
def main():
data_root = "/data/flower_photos" # get data root path
if not os.path.exists("./save_weights"):
os.makedirs("./save_weights")
batch_size = 8
epochs = 10
num_classes = 5
freeze_layers = True
initial_lr = 0.001
weight_decay = 1e-4
log_dir = "./logs/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
train_writer = tf.summary.create_file_writer(os.path.join(
log_dir, "train"))
val_writer = tf.summary.create_file_writer(os.path.join(log_dir, "val"))
# data generator with data augmentation
train_ds, val_ds = generate_ds(data_root,
batch_size=batch_size,
val_rate=0.2)
# create model
model = create_model(num_classes=num_classes, has_logits=False)
model.build((1, 224, 224, 3))
# 下载我提前转好的预训练权重
# 链接: https://pan.baidu.com/s/1ro-6bebc8zroYfupn-7jVQ 密码: s9d9
# load weights
pre_weights_path = './ViT-B_16.h5'
assert os.path.exists(pre_weights_path), "cannot find {}".format(
pre_weights_path)
model.load_weights(pre_weights_path, by_name=True, skip_mismatch=True)
# freeze bottom layers
if freeze_layers:
for layer in model.layers:
if "pre_logits" not in layer.name and "head" not in layer.name:
layer.trainable = False
else:
print("training {}".format(layer.name))
model.summary()
# custom learning rate curve
def scheduler(now_epoch):
end_lr_rate = 0.01 # end_lr = initial_lr * end_lr_rate
rate = ((1 + math.cos(now_epoch * math.pi / epochs)) /
2) * (1 - end_lr_rate) + end_lr_rate # cosine
new_lr = rate * initial_lr
# writing lr into tensorboard
with train_writer.as_default():
tf.summary.scalar('learning rate', data=new_lr, step=epoch)
return new_lr
# using keras low level api for training
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True)
optimizer = tf.keras.optimizers.SGD(learning_rate=initial_lr, momentum=0.9)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
name='train_accuracy')
val_loss = tf.keras.metrics.Mean(name='val_loss')
val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
name='val_accuracy')
@tf.function
def train_step(train_images, train_labels):
with tf.GradientTape() as tape:
output = model(train_images, training=True)
# cross entropy loss
ce_loss = loss_object(train_labels, output)
# l2 loss
matcher = re.compile(".*(bias|gamma|beta).*")
l2loss = weight_decay * tf.add_n([
tf.nn.l2_loss(v)
for v in model.trainable_variables if not matcher.match(v.name)
])
loss = ce_loss + l2loss
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(ce_loss)
train_accuracy(train_labels, output)
@tf.function
def val_step(val_images, val_labels):
output = model(val_images, training=False)
loss = loss_object(val_labels, output)
val_loss(loss)
val_accuracy(val_labels, output)
best_val_acc = 0.
for epoch in range(epochs):
train_loss.reset_states() # clear history info
train_accuracy.reset_states() # clear history info
val_loss.reset_states() # clear history info
val_accuracy.reset_states() # clear history info
# train
train_bar = tqdm(train_ds, file=sys.stdout)
for images, labels in train_bar:
train_step(images, labels)
# print train process
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(
epoch + 1, epochs, train_loss.result(),
train_accuracy.result())
# update learning rate
optimizer.learning_rate = scheduler(epoch)
# validate
val_bar = tqdm(val_ds, file=sys.stdout)
for images, labels in val_bar:
val_step(images, labels)
# print val process
val_bar.desc = "valid epoch[{}/{}] loss:{:.3f}, acc:{:.3f}".format(
epoch + 1, epochs, val_loss.result(), val_accuracy.result())
# writing training loss and acc
with train_writer.as_default():
tf.summary.scalar("loss", train_loss.result(), epoch)
tf.summary.scalar("accuracy", train_accuracy.result(), epoch)
# writing validation loss and acc
with val_writer.as_default():
tf.summary.scalar("loss", val_loss.result(), epoch)
tf.summary.scalar("accuracy", val_accuracy.result(), epoch)
# only save best weights
if val_accuracy.result() > best_val_acc:
best_val_acc = val_accuracy.result()
save_name = "./save_weights/model.ckpt"
model.save_weights(save_name, save_format="tf")