if opt.backbone == 'ResNet50':
model = PCB(opt)
elif opt.backbone == 'EfficientNet-B0':
model = PCB_Effi(opt)
if opt.LSTM:
# model_name = 'PCB-128_dim_cls'
# model = load_network(model, model_name)
model = PCB_Effi_LSTM(model, opt)
# model_name = 'LSTM'
# model = load_network(model, model_name)
if opt.GGNN:
# model_name = 'PCB-128_dim_cls'
# model = load_network(model, model_name)
model = PCB_Effi_GGNN(model, opt)
# model_name = 'LSTM'
# model = load_network(model, model_name)
print(model)
if opt.single_cls:
if opt.backbone == 'ResNet50':
ignored_params = list(map(id, model.model.fc.parameters()))
elif opt.backbone == 'EfficientNet-B0':
ignored_params = list(map(id, model.model._fc.parameters()))
ignored_params += (list(map(id, model.classifier.parameters())))
if opt.freeze_backbone:
ignored_params += (list(map(id, model.model.parameters())))
base_params = filter(
lambda p: id(p) not in ignored_params, model.parameters()
if opt.backbone == 'ResNet50':
model = PCB(opt)
elif opt.backbone == 'EfficientNet-B0':
model = PCB_Effi(opt)
if opt.LSTM:
model_name = 'test-pcb-ac'
model = load_network(model, model_name)
model = PCB_Effi_LSTM(model)
# model_name = 'LSTM'
# model = load_network(model, model_name)
if opt.GGNN:
model_name = 'test-pcb-ac'
model = load_network(model, model_name)
model = PCB_Effi_GGNN(model)
# model_name = 'GGNN'
# model = load_network(model, model_name)
print(model)
if opt.single_cls:
if opt.backbone == 'ResNet50':
ignored_params = list(map(id, model.model.fc.parameters()))
elif opt.backbone == 'EfficientNet-B0':
ignored_params = list(map(id, model.model._fc.parameters()))
ignored_params += (list(map(id, model.classifier.parameters())))
if opt.freeze_backbone:
ignored_params += (list(map(id, model.model.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
opt.nclasses = len(class_names)
if opt.PCB:
model = PCB_Effi(opt.nclasses)
if opt.LSTM:
# model_name = 'PCB-128_dim_cls'
# model = load_network(model, model_name)
model = PCB_Effi_LSTM(model, opt.train_backbone)
# model_name = 'LSTM'
# model = load_network(model, model_name)
if opt.GGNN:
model_name = 'PCB-128_dim_cls'
model = load_network(model, model_name)
model = PCB_Effi_GGNN(model, opt.train_backbone)
# model_name = 'LSTM'
# model = load_network(model, model_name)
print(model)
if not opt.multi_loss:
ignored_params = list(map(id, model.model._fc.parameters()))
ignored_params += (list(map(id, model.classifier.parameters())))
if not opt.train_backbone:
ignored_params += (list(map(id, model.model.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer = optim.SGD([{
'params': base_params,
'lr': 0.1 * opt.lr
opt.nclasses = len(class_names)
if opt.use_dense:
model = ft_net_dense(opt.nclasses, opt.droprate)
elif opt.use_NAS:
model = ft_net_NAS(opt.nclasses, opt.droprate)
else:
model = ft_net(opt.nclasses, opt.droprate, opt.stride)
if opt.PCB:
model = PCB_Effi(opt.nclasses)
if opt.PCB and (opt.LSTM or opt.GGNN):
model_name = 'PCB-128_dim_cls'
model = load_network(model, model_name)
model = PCB_Effi_LSTM(model) if opt.LSTM else PCB_Effi_GGNN(model)
# model_name = 'LSTM' or 'GGNN'
# model = load_network(model, model_name)
print(model)
if not opt.multi_loss:
ignored_params = list(map(id, model.model._fc.parameters()))
ignored_params += (list(map(id, model.classifier.parameters())) +
list(map(id, model.model.parameters())))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD([{
'params': base_params,
'lr': 0.1 * opt.lr
}, {
# Load Collected data Trained model
# --------
#
print('-------test-----------')
if opt.backbone == 'ResNet50':
model_structure = PCB(opt)
elif opt.backbone == 'EfficientNet-B0':
model_structure = PCB_Effi(opt)
if opt.LSTM:
model_structure = PCB_Effi_LSTM(model_structure)
if opt.GGNN:
model_structure = PCB_Effi_GGNN(model_structure)
model = load_network(model_structure)
# Remove the final fc layer and classifier layer
if not opt.LSTM and not opt.GGNN:
if opt.backbone == 'ResNet50':
model = PCB_test(model)
elif opt.backbone == 'EfficientNet-B0':
model = PCB_Effi_test(model)
elif opt.LSTM:
model = PCB_Effi_LSTM_test(model)
elif opt.GGNN:
model = PCB_Effi_GGNN_test(model)
else:
model.classifier.classifier = nn.Sequential()