#-------------------------------------------------------------------------------------#
input_shape = [600, 600, 3]
#----------------------------------------------------#
# 使用到的主干特征提取网络
# vgg或者resnet50
#----------------------------------------------------#
backbone = "resnet50"
model = FasterRCNN(NUM_CLASSES, backbone=backbone)
# #------------------------------------------------------#
# 权值文件请看README，百度网盘下载
#------------------------------------------------------#
model_path = 'model_data/voc_weights_resnet.pth'
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict() #OrderedDict :328
pretrained_dict = torch.load(model_path,
map_location=device) #OrderedDict :328
pretrained_dict = {
k: v
for k, v in pretrained_dict.items()
if np.shape(model_dict[k]) == np.shape(v)
}
model_dict.update(pretrained_dict) #更新模型的预训练参数
model.load_state_dict(model_dict)
print('Finished!')
net = model.train() #设置model为训练模式
if Cuda:
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
if __name__ == "__main__":
# 参数初始化
annotation_path = '2007_train.txt'
EPOCH_LENGTH = 2000
NUM_CLASSES = 20
IMAGE_SHAPE = [600,600,3]
BACKBONE = "resnet50"
model = FasterRCNN(NUM_CLASSES,backbone=BACKBONE).cuda()
#-------------------------------------------#
# 权值文件的下载请看README
#-------------------------------------------#
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load("model_data/voc_weights_resnet.pth", map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Finished!')
cudnn.benchmark = True
# 0.1用于验证，0.9用于训练
val_split = 0.1
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
np.random.seed(None)
args = parser.parse_args()
# 参数初始化
annotation_path = '2007_train.txt'
EPOCH_LENGTH = 2000
IMAGE_SHAPE = [600,600,3]
model = FasterRCNN(num_classes,backbone=args.backbone).cuda()
model_urls = {'resnet50':'https://github.com/you-bowen/tutorical_myDL/releases/download/1.0/frcnn_resnet50.pth',
'vgg':'none'}
#-------------------------------------------#
# 权值文件的下载请看README
#-------------------------------------------#
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load(args.pre_model) if args.pre_model else load_url(model_urls[args.backbone], map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v) and k[-19:] != 'num_batches_tracked'}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print(model.state_dict()['extractor.4.0.bn1.num_batches_tracked'])
print('Finished!')
cudnn.benchmark = True
# 0.1用于验证，0.9用于训练
val_split = args.val_split
with open(annotation_path) as f:
lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
