def main():
args = get_arguments()
myargs = [] # getopts(sys.argv)
now = datetime.datetime.now()
cwd = os.getcwd().replace('/tests', '')
if len(myargs) > 0:
if 'c' in myargs:
config_file = myargs['c']
else:
config_file = 'config/trainer_config.yml'
config = OmegaConf.load(os.path.join(cwd, config_file))['trainer']
config.cwd = str(cwd)
reproducibility(config)
print(f'pyTorch VERSION:{torch.__version__}', )
print(f'CUDA VERSION')
print(f'CUDNN VERSION:{torch.backends.cudnn.version()}')
print(f'Number CUDA Devices: {torch.cuda.device_count()}')
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
print(f'device: {device}')
training_generator, val_generator, test_generator, class_dict = select_dataset(
config)
n_classes = len(class_dict)
for loader in [training_generator, val_generator, test_generator]:
with torch.no_grad():
for batch_idx, (data, target) in enumerate(loader):
print(data.shape, target.shape)
# Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from utils.util import AttrDict, merge_cfg_from_args, get_arguments import os args = get_arguments() cfg = AttrDict() # 待预测图像所在路径 cfg.data_dir = os.path.join(args.example, "data", "test_images") # 待预测图像名称列表 cfg.data_list_file = os.path.join(args.example, "data", "test.txt") # 模型加载路径 cfg.model_path = os.path.join(args.example, "model") # 预测结果保存路径 cfg.vis_dir = os.path.join(args.example, "result") # 预测类别数 cfg.class_num = 2 # 均值, 图像预处理减去的均值 cfg.MEAN = 127.5, 127.5, 127.5
def main():
args = get_arguments()
myargs = [] # getopts(sys.argv)
now = datetime.datetime.now()
cwd = os.getcwd()
if len(myargs) > 0:
if 'c' in myargs:
config_file = myargs['c']
else:
config_file = 'config/trainer_config.yml'
config = OmegaConf.load(os.path.join(cwd, config_file))['trainer']
config.cwd = str(cwd)
reproducibility(config)
dt_string = now.strftime("%d_%m_%Y_%H.%M.%S")
cpkt_fol_name = os.path.join(
config.cwd,
f'checkpoints/model_{config.model.name}/dataset_{config.dataset.name}/date_{dt_string}'
)
log = Logger(path=cpkt_fol_name, name='LOG').get_logger()
best_pred_loss = 1000.0
log.info(f"Checkpoint folder {cpkt_fol_name}")
log.info(f"date and time = {dt_string}")
log.info(f'pyTorch VERSION:{torch.__version__}', )
log.info(f'CUDA VERSION')
log.info(f'CUDNN VERSION:{torch.backends.cudnn.version()}')
log.info(f'Number CUDA Devices: {torch.cuda.device_count()}')
if args.tensorboard:
writer_path = os.path.join(cpkt_fol_name + 'runs/')
writer = SummaryWriter(writer_path + util.datestr())
else:
writer = None
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
log.info(f'device: {device}')
training_generator, val_generator, test_generator, class_dict = select_dataset(
config)
n_classes = len(class_dict)
model = select_model(config, n_classes)
log.info(f"{model}")
if (config.load):
pth_file, _ = load_checkpoint(config.pretrained_cpkt,
model,
strict=True,
load_seperate_layers=False)
else:
pth_file = None
if (config.cuda and use_cuda):
if torch.cuda.device_count() > 1:
log.info(f"Let's use {torch.cuda.device_count()} GPUs!")
model = torch.nn.DataParallel(model)
model.to(device)
optimizer, scheduler = select_optimizer(model, config['model'], None)
log.info(f'{model}')
log.info(f"Checkpoint Folder {cpkt_fol_name} ")
shutil.copy(os.path.join(config.cwd, config_file), cpkt_fol_name)
trainer = Trainer(config,
model=model,
optimizer=optimizer,
data_loader=training_generator,
writer=writer,
logger=log,
valid_data_loader=val_generator,
test_data_loader=test_generator,
class_dict=class_dict,
lr_scheduler=scheduler,
checkpoint_dir=cpkt_fol_name)
trainer.train()
