config = Config.fromfile(args.config)
config.workspace = os.path.join(config.workspace_dir, config.exp_name)
config.checkpoint = args.resume_from
config.visualization = config.visualization
if not os.path.exists(config.workspace):
os.makedirs(config.workspace)
logger = setup_logger(os.path.join(config.workspace, 'test_log'))
logger.info(config.print())
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.kernel_num,
scale=config.scale)
num_gpus = torch.cuda.device_count()
device = torch.device("cuda:0")
model = nn.DataParallel(model)
model = model.to(device)
state = torch.load(config.checkpoint)
model.load_state_dict(state['state_dict'])
logger.info('test epoch {}'.format(state['epoch']))
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
start_epoch = load_checkpoint(config.checkpoint, model, logger, device,
optimizer)
for i in range(60, 100, 2):
thre = i * 0.01
print(thre)
eval(model, config, device, thre)
def main():
if config.output_dir is None:
config.output_dir = 'output'
if config.restart_training:
shutil.rmtree(config.output_dir, ignore_errors=True)
if not os.path.exists(config.output_dir):
os.makedirs(config.output_dir)
logger = setup_logger(os.path.join(config.output_dir, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = MyDataset(config.trainroot,
data_shape=config.data_shape,
n=config.n,
m=config.m,
transform=transforms.ToTensor())
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers),
drop_last=True)
writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.n,
scale=config.scale)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
## loading the pretrained weights from drive
state_dict = torch.load(config.pretrained_path)
model.load_state_dict(state_dict)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
# dummy_input = torch.autograd.Variable(torch.Tensor(1, 3, 600, 800).to(device))
# writer.add_graph(models=models, input_to_model=dummy_input)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
# optimizer = torch.optim.SGD(models.parameters(), lr=config.lr, momentum=0.99)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device)
start_epoch += 1
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
try:
for epoch in range(start_epoch, config.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler,
train_loader, device, criterion,
epoch, all_step, writer, logger)
logger.info(
'[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
# net_save_path = '{}/PSENet_{}_loss{:.6f}.pth'.format(config.output_dir, epoch,
# train_loss)
# save_checkpoint(net_save_path, models, optimizer, epoch, logger)
state_dict = model.state_dict()
# replace the weight file
filename = '{}/PSENet_resnet50.pth'.format(config.output_dir)
if os.path.exists(filename):
os.unlink(filename)
torch.save(state_dict, filename)
writer.close()
except KeyboardInterrupt:
filename = '{}/PSENet_resnet50.pth'.format(config.output_dir)
if os.path.exists(filename):
os.unlink(filename)
torch.save(state_dict, filename)
save_checkpoint('{}/final.pth'.format(config.output_dir), model,
optimizer, epoch, logger)