if __name__ == '__main__':
import argparse
# from efficientnet_pytorch import EfficientNet
parser = argparse.ArgumentParser(
description='Convert TF model to PyTorch model and save for easier future loading')
parser.add_argument('--model_name', type=str, default='efficientnet-b0',
help='efficientnet-b{N}, where N is an integer 0 <= N <= 7')
parser.add_argument('--pth_file', type=str, default='efficientnet-b0.pth',
help='input PyTorch model file name')
args = parser.parse_args()
# Build model
model = EfficientNet.from_name(args.model_name)
pretrained_weights = torch.load(args.pth_file)
# model.load_state_dict(pretrained_weights)#error,key mismatched
# print(type(pretrained_weights),dir(pretrained_weights))#<class 'collections.OrderedDict'>
# for key in pretrained_weights.keys():
# print(key)
del pretrained_weights['_fc.weight']#delete unuseful weights
del pretrained_weights['_fc.bias']
# for key in pretrained_weights.keys():
# print(key)
model.load_state_dict(pretrained_weights)
# from torchsummary import summary
# summary(model.cuda(), input_size=(3, 320, 320))
print('*' * 30)
train_loader = get_loader(osp.join(args.data_path, 'images'),
osp.join(args.data_path, 'list_attr_celeba.txt'),
crop_size=178,
image_size=360,
batch_size=args.batch_size,
mode='all',
num_workers=args.num_workers)
print('*' * 30)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('Running on device:', device)
print('*' * 30)
print('Model Architecture')
model = EfficientNet.from_name('efficientnet-b4')
model = model.to(device)
model_params = torchutils.get_model_param_count(model)
print(model)
print('Total model parameters:', model_params)
print('*' * 30)
# criterion = nn.BCEWithLogitsLoss()
criterion = nn.MultiLabelSoftMarginLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
threshold=0.005,