200 lines
6.0 KiB
Python
200 lines
6.0 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
#
|
|
# Copyright (c) 2021, Pl@ntNet
|
|
# All rights reserved.
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
#
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
#
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
import os
|
|
from tqdm import tqdm
|
|
import pickle
|
|
import argparse
|
|
import time
|
|
import torch
|
|
from torch.optim import SGD
|
|
from torch.nn import CrossEntropyLoss
|
|
|
|
from utils import set_seed, load_model, save, get_model, update_optimizer, get_data
|
|
from epoch import train_epoch, val_epoch, test_epoch
|
|
from cli import add_all_parsers
|
|
|
|
|
|
def train(args):
|
|
set_seed(args, use_gpu=torch.cuda.is_available())
|
|
train_loader, val_loader, test_loader, dataset_attributes = get_data(
|
|
args.root,
|
|
args.image_size,
|
|
args.crop_size,
|
|
args.batch_size,
|
|
args.num_workers,
|
|
args.pretrained,
|
|
)
|
|
|
|
model = get_model(args, n_classes=dataset_attributes["n_classes"])
|
|
criteria = CrossEntropyLoss()
|
|
|
|
if args.use_gpu:
|
|
print("USING GPU")
|
|
torch.cuda.set_device(0)
|
|
model.cuda()
|
|
criteria.cuda()
|
|
|
|
optimizer = SGD(
|
|
model.parameters(),
|
|
lr=args.lr,
|
|
momentum=0.9,
|
|
weight_decay=args.mu,
|
|
nesterov=True,
|
|
)
|
|
|
|
# Containers for storing metrics over epochs
|
|
loss_train, acc_train, topk_acc_train = [], [], []
|
|
loss_val, acc_val, topk_acc_val, avgk_acc_val, class_acc_val = [], [], [], [], []
|
|
|
|
save_name = args.save_name_xp.strip()
|
|
save_dir = os.path.join(os.getcwd(), "results", save_name)
|
|
if not os.path.exists(save_dir):
|
|
os.makedirs(save_dir)
|
|
|
|
print("args.k : ", args.k)
|
|
|
|
lmbda_best_acc = None
|
|
best_val_acc = float("-inf")
|
|
|
|
for epoch in tqdm(range(args.n_epochs), desc="epoch", position=0):
|
|
t = time.time()
|
|
optimizer = update_optimizer(
|
|
optimizer, lr_schedule=args.epoch_decay, epoch=epoch
|
|
)
|
|
|
|
loss_epoch_train, acc_epoch_train, topk_acc_epoch_train = train_epoch(
|
|
model,
|
|
optimizer,
|
|
train_loader,
|
|
criteria,
|
|
loss_train,
|
|
acc_train,
|
|
topk_acc_train,
|
|
args.k,
|
|
dataset_attributes["n_train"],
|
|
args.use_gpu,
|
|
)
|
|
|
|
(
|
|
loss_epoch_val,
|
|
acc_epoch_val,
|
|
topk_acc_epoch_val,
|
|
avgk_acc_epoch_val,
|
|
lmbda_val,
|
|
) = val_epoch(
|
|
model,
|
|
val_loader,
|
|
criteria,
|
|
loss_val,
|
|
acc_val,
|
|
topk_acc_val,
|
|
avgk_acc_val,
|
|
class_acc_val,
|
|
args.k,
|
|
dataset_attributes,
|
|
args.use_gpu,
|
|
)
|
|
|
|
# save model at every epoch
|
|
save(
|
|
model, optimizer, epoch, os.path.join(save_dir, save_name + "_weights.tar")
|
|
)
|
|
|
|
# save model with best val accuracy
|
|
if acc_epoch_val > best_val_acc:
|
|
best_val_acc = acc_epoch_val
|
|
lmbda_best_acc = lmbda_val
|
|
save(
|
|
model,
|
|
optimizer,
|
|
epoch,
|
|
os.path.join(save_dir, save_name + "_weights_best_acc.tar"),
|
|
)
|
|
|
|
print()
|
|
print(f"epoch {epoch} took {time.time()-t:.2f}")
|
|
print(f"loss_train : {loss_epoch_train}")
|
|
print(f"loss_val : {loss_epoch_val}")
|
|
print(
|
|
f"acc_train : {acc_epoch_train} / topk_acc_train : {topk_acc_epoch_train}"
|
|
)
|
|
print(
|
|
f"acc_val : {acc_epoch_val} / topk_acc_val : {topk_acc_epoch_val} / "
|
|
f"avgk_acc_val : {avgk_acc_epoch_val}"
|
|
)
|
|
|
|
# load weights corresponding to best val accuracy and evaluate on test
|
|
load_model(
|
|
model, os.path.join(save_dir, save_name + "_weights_best_acc.tar"), args.use_gpu
|
|
)
|
|
(
|
|
loss_test_ba,
|
|
acc_test_ba,
|
|
topk_acc_test_ba,
|
|
avgk_acc_test_ba,
|
|
class_acc_test,
|
|
) = test_epoch(
|
|
model,
|
|
test_loader,
|
|
criteria,
|
|
args.k,
|
|
lmbda_best_acc,
|
|
args.use_gpu,
|
|
dataset_attributes,
|
|
)
|
|
|
|
# Save the results as a dictionary and save it as a pickle file in desired location
|
|
|
|
results = {
|
|
"loss_train": loss_train,
|
|
"acc_train": acc_train,
|
|
"topk_acc_train": topk_acc_train,
|
|
"loss_val": loss_val,
|
|
"acc_val": acc_val,
|
|
"topk_acc_val": topk_acc_val,
|
|
"class_acc_val": class_acc_val,
|
|
"avgk_acc_val": avgk_acc_val,
|
|
"test_results": {
|
|
"loss": loss_test_ba,
|
|
"accuracy": acc_test_ba,
|
|
"topk_accuracy": topk_acc_test_ba,
|
|
"avgk_accuracy": avgk_acc_test_ba,
|
|
"class_acc_dict": class_acc_test,
|
|
},
|
|
"params": args.__dict__,
|
|
}
|
|
|
|
with open(os.path.join(save_dir, save_name + ".pkl"), "wb") as f:
|
|
pickle.dump(results, f)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
add_all_parsers(parser)
|
|
args = parser.parse_args()
|
|
train(args)
|