#-------------------------------------------------------------------------------------# 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)