def main():
#加载数据
trainloader, testloader = read_dataset(input_size, batch_size, root, set)
#定义模型
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model).cuda()
#设置训练参数
criterion_bp = nn.CrossEntropyLoss()
criterion_ls = nn.BCELoss() # for mura
parameters = model.parameters()
#加载checkpoint
save_path = os.path.join(model_path, model_name)
if os.path.exists(save_path):
start_epoch, lr = auto_load_resume(model, save_path, status='train')
assert start_epoch < end_epoch
else:
os.makedirs(save_path)
start_epoch = 0
lr = init_lr
# define optimizers
optimizer = torch.optim.SGD(parameters,
lr=lr,
momentum=0.9,
weight_decay=weight_decay)
model = model.cuda() # 部署在GPU
scheduler = MultiStepLR(optimizer,
milestones=lr_milestones,
gamma=lr_decay_rate)
# 保存config参数信息
time_str = time.strftime("%Y%m%d-%H%M%S")
shutil.copy('./config.py',
os.path.join(save_path, "{}config.py".format(time_str)))
# 开始训练
train(model=model,
trainloader=trainloader,
testloader=testloader,
criterion_bp=criterion_bp,
criterion_ls=criterion_ls,
optimizer=optimizer,
scheduler=scheduler,
save_path=save_path,
start_epoch=start_epoch,
end_epoch=end_epoch,
save_interval=save_interval)
def main():
trainset, _, testset, _ = read_dataset(input_size, batch_size, root, set)
# image will be resize to the input_size
# batch size means the number of images the nn process before updating the weight and biases
# root is the root to the dataset
# set is the dataset name (change in config)
# Load checkpoint from a fold number
save_path = os.path.join(model_path, model_name)
if os.path.exists(save_path):
load_model_from_path = os.path.join(save_path,
f'fold_{start_from_fold}')
if not os.path.exists(load_model_from_path):
os.makedirs(load_model_from_path)
# Create model
if (multitask):
model = MainNetMultitask(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
else:
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
start_epoch, lr, patience_counter = auto_load_resume(
model, load_model_from_path, status='train')
print(f'Patience counter starting from: {patience_counter}')
assert start_epoch < end_epoch, 'end of fold reached, please increment start_from_fold'
assert start_from_fold < num_folds
assert patience_counter <= patience
else:
os.makedirs(save_path)
start_epoch = 0
lr = init_lr
patience_counter = 0
# save config
time_str = time.strftime("%Y%m%d-%H%M%S")
shutil.copy('./config.py',
os.path.join(save_path, "{}config.py".format(time_str)))
print(
'\nSplitting trainset into train and val sets (80:20)\nNote testset is loaded but unused, and will not be used unless test.py is run.'
)
# NOTE: split train into train/val set; but for consistency of code we'll leave the variable names as 'test' instead of 'val'
# split train set into train/val 80:20
X_train = []
y_train = []
for i, j in trainset.train_img_label:
X_train.append(i)
y_train.append(j)
# convert lists into numpy arrays
X_train = np.array(X_train)
y_train = np.array(y_train)
skf = StratifiedKFold(n_splits=num_folds, random_state=seed, shuffle=True)
if (multitask):
# placeholder y_train with only the first element of the y tuples for when multitask learning is done to prevent stratified kfolds bug
y_train_temp = np.array([i[0] for i in y_train])
for fold, (train_index,
val_index) in enumerate(skf.split(X_train, y_train_temp)):
print(f'Multitask: {multitask}')
print(f'Random Augmentation: {rand_aug}')
if (rand_aug):
print(f'N:{N} M: {M}')
print(f'\n=============== Fold {fold+1} ==================')
if (fold + 1 < start_from_fold):
print('Skipping this fold...\n')
continue
# Prepare save_path for the fold
save_path_fold = os.path.join(save_path, str(f'fold_{fold+1}'))
# Split trainloader into train and val loaders
X_train_fold = X_train[train_index]
X_val_fold = X_train[val_index]
y_train_fold = y_train[train_index]
y_val_fold = y_train[val_index]
# Zip back the X and y values
train_img_label_fold = list(zip(X_train_fold, y_train_fold))
val_img_label_fold = list(zip(X_val_fold, y_val_fold))
# Hijack the original trainset with the X and y for the particular fold
trainset_fold = trainset
trainset_fold.train_img_label = train_img_label_fold
valset_fold = testset
valset_fold.test_img_label = val_img_label_fold # variable name kept as test_img_label for code consistency
print(f'Size of trainset: {len(trainset_fold)}')
print(f'Size of valset: {len(valset_fold)}')
# Recreate DataLoaders with train and val sets
trainloader_fold = torch.utils.data.DataLoader(
trainset_fold,
batch_size=batch_size,
shuffle=True,
num_workers=8,
drop_last=False)
valloader_fold = torch.utils.data.DataLoader(valset_fold,
batch_size=batch_size,
shuffle=False,
num_workers=8,
drop_last=False)
# Create model
model = MainNetMultitask(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
criterion = nn.CrossEntropyLoss()
# Define optimizers
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters,
lr=lr,
momentum=0.9,
weight_decay=weight_decay)
model = model.cuda()
scheduler = MultiStepLR(optimizer,
milestones=lr_milestones,
gamma=lr_decay_rate)
train_multitask(model=model,
trainloader=trainloader_fold,
testloader=valloader_fold,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
save_path=save_path_fold,
start_epoch=start_epoch,
end_epoch=end_epoch,
patience_counter=patience_counter,
save_interval=save_interval)
start_epoch = 0 # refresh start_epoch for next fold
# Clear model and release GPU memory
del model
torch.cuda.empty_cache()
print(
f'\n=============== End of fold {fold+1} ==================\n')
else:
for fold, (train_index,
val_index) in enumerate(skf.split(X_train, y_train)):
print(f'Multitask: {multitask}')
print(f'Random Augmentation: {rand_aug}')
if (rand_aug):
print(f'N:{N} M: {M}')
print(f'\n=============== Fold {fold+1} ==================')
if (fold + 1 < start_from_fold):
print('Skipping this fold...\n')
continue
# Prepare save_path for the fold
save_path_fold = os.path.join(save_path, str(f'fold_{fold+1}'))
# Split trainloader into train and val loaders
X_train_fold = X_train[train_index]
X_val_fold = X_train[val_index]
y_train_fold = y_train[train_index]
y_val_fold = y_train[val_index]
# Zip back the X and y values
train_img_label_fold = list(zip(X_train_fold, y_train_fold))
val_img_label_fold = list(zip(X_val_fold, y_val_fold))
# Hijack the original trainset with the X and y for the particular fold
trainset_fold = trainset
trainset_fold.train_img_label = train_img_label_fold
valset_fold = testset
valset_fold.test_img_label = val_img_label_fold # variable name kept as test_img_label for code consistency
print(f'Size of trainset: {len(trainset_fold)}')
print(f'Size of valset: {len(valset_fold)}')
# Recreate DataLoaders with train and val sets
trainloader_fold = torch.utils.data.DataLoader(
trainset_fold,
batch_size=batch_size,
shuffle=True,
num_workers=8,
drop_last=False)
valloader_fold = torch.utils.data.DataLoader(valset_fold,
batch_size=batch_size,
shuffle=False,
num_workers=8,
drop_last=False)
# Create model
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
criterion = nn.CrossEntropyLoss()
# Define optimizers
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters,
lr=lr,
momentum=0.9,
weight_decay=weight_decay)
model = model.cuda()
scheduler = MultiStepLR(optimizer,
milestones=lr_milestones,
gamma=lr_decay_rate)
train(model=model,
trainloader=trainloader_fold,
testloader=valloader_fold,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
save_path=save_path_fold,
start_epoch=start_epoch,
end_epoch=end_epoch,
patience_counter=patience_counter,
save_interval=save_interval)
start_epoch = 0 # refresh start_epoch for next fold
# Clear model and release GPU memory
del model
torch.cuda.empty_cache()
print(
f'\n=============== End of fold {fold+1} ==================\n')
set = 'Aircraft'
if set == 'CUB':
root = './datasets/CUB_200_2011' # dataset path
# model path
pth_path = "./models/cub_epoch144.pth"
num_classes = 200
elif set == 'Aircraft':
root = './datasets/FGVC-aircraft' # dataset path
# model path
pth_path = "./models/air_epoch146.pth"
num_classes = 100
batch_size = 10
#load dataset
_, testloader = read_dataset(input_size, batch_size, root, set)
# 定义模型
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
model = model.to(DEVICE)
criterion = nn.CrossEntropyLoss()
#加载checkpoint
if os.path.exists(pth_path):
epoch = auto_load_resume(model, pth_path, status='test')
else:
sys.exit('There is not a pth exist.')
def main():
_, trainloader, _, testloader = read_dataset(input_size, batch_size, root,
set)
# image will be resize to the input_size
# batch size means the number of images the nn process before updating the weight and biases
# root is the root to the dataset
# set is the dataset name (change in config)
# Load checkpoint
save_path = os.path.join(model_path, model_name)
if os.path.exists(save_path):
load_model_from_path = save_path
if not os.path.exists(load_model_from_path):
os.makedirs(load_model_from_path)
# Create model
if (multitask):
model = MainNetMultitask(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
else:
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
start_epoch, lr, patience_counter = auto_load_resume(
model, load_model_from_path, status='train')
print(f'Patience counter starting from: {patience_counter}')
assert start_epoch < end_epoch, 'maximum number of epochs reached'
assert patience_counter <= patience
else:
os.makedirs(save_path)
start_epoch = 0
lr = init_lr
patience_counter = 0
if (multitask):
print(f'Multitask: {multitask}')
print(f'Random Augmentation: {rand_aug}')
if (rand_aug):
print(f'N:{N} M: {M}')
# Create model
model = MainNetMultitask(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
criterion = nn.CrossEntropyLoss()
# Define optimizers
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters,
lr=lr,
momentum=0.9,
weight_decay=weight_decay)
model = model.cuda()
scheduler = MultiStepLR(optimizer,
milestones=lr_milestones,
gamma=lr_decay_rate)
train_multitask(model=model,
trainloader=trainloader,
testloader=testloader,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
save_path=save_path,
start_epoch=start_epoch,
end_epoch=end_epoch,
patience_counter=patience_counter,
save_interval=save_interval)
else:
print(f'Multitask: {multitask}')
print(f'Random Augmentation: {rand_aug}')
if (rand_aug):
print(f'N:{N} M: {M}')
# Create model
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
criterion = nn.CrossEntropyLoss()
# Define optimizers
parameters = model.parameters()
optimizer = torch.optim.SGD(parameters,
lr=lr,
momentum=0.9,
weight_decay=weight_decay)
model = model.cuda()
scheduler = MultiStepLR(optimizer,
milestones=lr_milestones,
gamma=lr_decay_rate)
train(model=model,
trainloader=trainloader,
testloader=testloader,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
save_path=save_path,
start_epoch=start_epoch,
end_epoch=end_epoch,
patience_counter=patience_counter,
save_interval=save_interval)
def main():
CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda" if CUDA else "cpu")
set = 'CompCars' # ensure dataset is set properly
if set == 'CUB':
root = './datasets/CUB_200_2011' # dataset path
# model path
pth_path = "./models/cub_epoch144.pth"
num_classes = 200
elif set == 'CompCars':
root = './datasets/CompCars' # dataset path
# model path
pth_path = "./models/car_model_randaug_epoch45.pth" # remember to change as per saved model's name
num_classes = 431
if (multitask):
num_classes = (431, 75)
assert isinstance(
num_classes, tuple
), "Multitask mode is enabled but num_classes is an integer; please pass in a tuple for multiple predictions."
assert num_classes[
0] == 431, "Please put the number of car models (431) as the first element of the tuple"
elif set == 'Aircraft':
root = './datasets/FGVC-aircraft' # dataset path
# model path
pth_path = "./models/air_epoch146.pth"
num_classes = 100
batch_size = 10
# load test dataset
_, _, _, testloader = read_dataset(input_size, batch_size, root, set)
# create model
if (multitask):
model = MainNetMultitask(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
else:
model = MainNet(proposalN=proposalN,
num_classes=num_classes,
channels=channels)
model = model.to(DEVICE)
criterion = nn.CrossEntropyLoss()
# Load checkpoint model
if os.path.exists(pth_path):
epoch = auto_load_resume(model, pth_path, status='test')
else:
sys.exit('No saved model checkpoint detected.')
print('Testing on saved model...')
if (multitask):
print('MULTITASK TEST...')
object_correct_1 = 0
object_correct_2 = 0
model.eval()
with torch.no_grad():
for i, data in enumerate(tqdm(testloader)):
if set == 'CUB':
x, y, boxes, _ = data
else:
x, y = data
x = x.to(DEVICE)
y_1, y_2 = y[0], y[1]
y_1 = y_1.to(DEVICE)
y_2 = y_2.to(DEVICE)
local_logits_1, local_logits_2, local_imgs = model(
x, epoch, i, 'test', DEVICE)[-3:]
# local
pred_1 = local_logits_1.max(1, keepdim=True)[1]
pred_2 = local_logits_2.max(1, keepdim=True)[1]
object_correct_1 += pred_1.eq(y_1.view_as(pred_1)).sum().item()
object_correct_2 += pred_2.eq(y_2.view_as(pred_2)).sum().item()
print('\nObject branch accuracy for task 1: {}/{} ({:.2f}%)\n'.
format(object_correct_1, len(testloader.dataset),
100. * object_correct_1 / len(testloader.dataset)))
print('\nObject branch accuracy for task 2: {}/{} ({:.2f}%)\n'.
format(object_correct_2, len(testloader.dataset),
100. * object_correct_2 / len(testloader.dataset)))
else:
object_correct = 0
model.eval()
with torch.no_grad():
for i, data in enumerate(tqdm(testloader)):
if set == 'CUB':
x, y, boxes, _ = data
else:
x, y = data
x = x.to(DEVICE)
y = y.to(DEVICE)
local_logits, local_imgs = model(x, epoch, i, 'test',
DEVICE)[-2:]
# local
pred = local_logits.max(1, keepdim=True)[1]
object_correct += pred.eq(y.view_as(pred)).sum().item()
print('\nObject branch accuracy: {}/{} ({:.2f}%)\n'.format(
object_correct, len(testloader.dataset),
100. * object_correct / len(testloader.dataset)))