def __init__(self): super(ToyPredictor, self).__init__() self.device = 'cuda' if torch.cuda.is_available() else 'cpu' self.is_ensamble = True if not self.is_ensamble: # model, _, *_ = model_selection('se_resnext101_32x4d', num_out_classes=2, dropout=0.5) model = get_efficientnet(model_name='efficientnet-b0', num_classes=2, pretrained=False) model_path = './weight/output_my_aug/efn-b0_LS_27_loss_0.2205.pth' model.load_state_dict( torch.load(model_path, map_location=self.device)) print('Load model in:', model_path) self.model = model.to(self.device) else: self.models = self.load_models( model_names=[ 'efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2' ], model_paths=[ './weight/output_my_aug/efn-b0_LS_27_loss_0.2205.pth', './weight/output_my_aug/efn-b1_LS_6_loss_0.1756.pth', './weight/output_my_aug/efn-b2_LS_12_loss_0.1728.pth' ]) self.mtcnn = MTCNN(margin=14, keep_all=True, factor=0.6, device=self.device).eval()
def load_models(self, model_names, model_paths): models = [] for i in range(len(model_names)): model = get_efficientnet(model_name=model_names[i], num_classes=2, pretrained=False) model_path = model_paths[i] model.load_state_dict( torch.load(model_path, map_location=self.device)) print('Load model ', i, 'in:', model_path) model.to(self.device) models.append(model) return models
if __name__ == '__main__': LOG_FREQ = 50 batch_size = 128 test_batch_size = 128 device_id = 0 lr = 1e-3 epoch_start = 1 num_epochs = epoch_start + 50 model_name = 'efficientnet-b1' writeFile = '/data1/cby/temp/output_my_aug/logs/' + model_name store_name = '/data1/cby/temp/output_my_aug/weights/' + model_name if not os.path.isdir(store_name): os.makedirs(store_name) model_path = None # model_path = '/data1/cby/temp/output_my_aug/weights/efficientnet-b1/efn-b1_LS_9_loss_0.1610.pth' model = get_efficientnet(model_name=model_name) if model_path is not None: # model = torch.load(model_path) model.load_state_dict(torch.load(model_path, map_location='cpu')) print('Model found in {}'.format(model_path)) else: print('No model found, initializing random model.') model = model.cuda(device_id) train_logger = Logger(model_name=writeFile, header=['epoch', 'loss', 'acc', 'lr']) # criterion = nn.CrossEntropyLoss() criterion = LabelSmoothing(smoothing=0.05).cuda(device_id) # optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9) # optimizer = optim.Adam(model.parameters(), lr=lr) optimizer = optim.AdamW(model.parameters(), lr=lr)