def load_model(): global file_num file_num = 0 global class_model class_model = feature_extractor() model_path = './class_model.pki' tmp = torch.load(model_path, map_location={'cuda:0': 'cpu'}) class_model.load_state_dict(tmp) class_model.eval() del tmp global object_model object_model = SSD(depth=50, width=1) model_path = './ssd_patch.pki' tmp = torch.load(model_path, map_location={'cuda:0': 'cpu'}) object_model.load_state_dict(tmp) object_model.eval()
def initialize_net() -> None: global ssd_net # if already defined, return it if ssd_net is not None: print('use cached ssd_net') return ssd_net # get device ( cpu / gpu ) to be used use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") print(f'device : {device}') ssd_cfg = { 'num_classes': num_classes, # number of classes including background class 'input_size': Parameters.IMG_SIZE, 'bbox_aspect_num': Parameters.BBOX_ASPECT_NUM, 'feature_maps': Parameters.FEATURE_MAPS, 'steps': Parameters.STEPS, 'min_sizes': Parameters.MIN_SIZES, 'max_sizes': Parameters.MAX_SIZES, 'aspect_ratios': Parameters.ASPECT_RATIOS, 'conf_thresh': Parameters.CONF_THRESHOLD, 'top_k': Parameters.TOP_K, 'nms_thresh': Parameters.NMS_THRESHOLD } print(f'initializing ssd with : {ssd_cfg}') ssd_net = SSD(phase="inference", cfg=ssd_cfg) # load weight created in training weight_file_path = os.path.join(Parameters.ABEJA_TRAINING_RESULT_DIR, 'model.pth') print(f'weight_file_path : {weight_file_path}') # cf. https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-on-gpu-load-on-gpu weight = torch.load(weight_file_path, map_location=device) ssd_net.load_state_dict(weight) ssd_net = ssd_net.to(device) ssd_net.eval() return ssd_net
def train(train_config): logger = Logger(HOME+'/log', train_config.basenet) if train_config.dataset_name == 'VOC': cfg = voc_config dataset = VOCDataset(DATA_DIR, transform=SSDAugmentation( cfg['min_dim'], MEANS)) elif train_config.dataset_name == 'COCO': cfg = coco_config dataset = COCODataset(DATA_DIR, transform=SSDAugmentation( cfg['min_dim'], MEANS)) if train_config.visdom: import visdom viz = visdom.Visdom() ssd_net = SSD('train', train_config.basenet, cfg['min_dim'], cfg['num_classes'], with_fpn=train_config.with_fpn) net = ssd_net if train_config.cuda: net = nn.DataParallel(ssd_net) cudnn.benchmark = True if train_config.resume: logger('Loading {} ...'.format(train_config.resume)) load_weights = torch.load( train_config.resume, map_location=lambda storage, loc: storage) ssd_net.load_state_dict(load_weights) if train_config.cuda: net = net.cuda() if not train_config.resume: logger('Initializing weights ...') ssd_net.topnet.apply(weights_init) ssd_net.loc_layers.apply(weights_init) ssd_net.conf_layers.apply(weights_init) optimizer = optim.Adam(net.parameters(), lr=train_config.lr, weight_decay=train_config.weight_decay) criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, False, train_config.cuda) net.train() # loss counters loc_loss = 0 conf_loss = 0 epoch = 0 logger('Loading the dataset...') epoch_size = len(dataset) // train_config.batch_size logger('Training SSD on:{}'.format(dataset.name)) # logger('using the specified args:') step_index = 0 if train_config.visdom: vis_title = 'SSD.PyTorch on ' + dataset.name vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss'] iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend) epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend) data_loader = data.DataLoader(dataset, train_config.batch_size, num_workers=train_config.num_workers, shuffle=True, collate_fn=detection_collate, pin_memory=True) # create batch iterator batch_iterator = iter(data_loader) t0 = time.time() for iteration in range(train_config.start_iter, cfg['max_iter']): if train_config.visdom and iteration != 0 and (iteration % epoch_size == 0): update_vis_plot(epoch, loc_loss.item(), conf_loss.item(), epoch_plot, None, 'append', epoch_size) logger('epoch = {} : loss = {}, loc_loss = {}, conf_loss = {}'.format( epoch, loc_loss + conf_loss, loc_loss, conf_loss)) # reset epoch loss counters loc_loss = 0 conf_loss = 0 epoch += 1 if iteration in cfg['lr_steps']: step_index += 1 adjust_learning_rate(optimizer, train_config.lr, train_config.gamma, step_index) # load train data images, targets = next(batch_iterator) if iteration//epoch_size > 0 and iteration % epoch_size == 0: batch_iterator = iter(data_loader) print(iteration) if train_config.cuda: images = images.cuda() targets = [ann.cuda()for ann in targets] # else: # images=torch.tensor(images) # targets=torch.tensor(targets) # forward out = net(images) # backprop optimizer.zero_grad() loss_l, loss_c = criterion(out, targets) loss = loss_l + loss_c loss.backward() optimizer.step() t1 = time.time() if train_config.visdom: loc_loss += loss_l.item() conf_loss += loss_c.item() if iteration % 50 == 0: t1 = time.time() logger('timer: %.4f sec. || ' % (t1 - t0)+'iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()) + ' || loc_loss: %.4f ||' % (loss_l.item()) + ' || conf_loss: %.4f ||' % (loss_c.item())) t0 = time.time() if train_config.visdom: update_vis_plot(iteration, loss_l.item(), loss_c.item(), iter_plot, epoch_plot, 'append') if iteration != 0 and iteration % 5000 == 0: logger('Saving state, iter:%d' % iteration) torch.save(ssd_net.state_dict(), train_config.save_folder + 'ssd224_VOC_' + repr(iteration) + '.pth') torch.save(ssd_net.state_dict(), train_config.save_folder + 'ssd224_VOC.pth')
'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]], } #SSDネットワークモデル net = SSD(phase="inference", cfg=ssd_cfg) #device = torch.device('cpu') #SSDの学習済み重みを設定 #net_weights = torch.load('./weights/ssd300_600.pth', map_location=device) #net_weights = torch.load('./weights/ssd300_50.pth', map_location={'cuda0': 'cpu'}) net_weights = torch.load('./weights/ssd300_1.pth', map_location={'cuda0': 'cpu'}) #net_weights = torch.load('./weights/ssd300_mAP_77.43_v2.pth', map_location={'cuda0': 'cpu'}) ###net_weights = torch.load('./weights/ssd300_mAP_77.43_v2.pth', map_location=device) net.load_state_dict(net_weights) print('ネットワーク設定完了:学習済みの重みをロードしました') #画像読み込み image_file_path = "./data/cowboy-757575_640.jpg" img = cv2.imread(image_file_path) #[高さ][幅][色BGR] height, width, chanels = img.shape #画像サイズを取得 #元画像の表示 plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) plt.show() #前処理クラスの作成 color_mean = (104, 117, 123) #(BGR)の色の平均値 input_size = 300 #画像のinputサイズを300×300にする