# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torch.optim as optim
from tensorboardX import SummaryWriter
import numpy as np
import time
import datetime
import argparse
import os
import os.path as osp
from rs_dataset import RSDataset
from get_logger import get_logger
from res_network import Resnet18, Resnet34, Resnet101, Densenet121, SEResNext50
def parse_args():
parse = argparse.ArgumentParser()
parse.add_argument('--epoch', type=int, default=15)
parse.add_argument('--schedule_step', type=int, default=4)
parse.add_argument('--batch_size', type=int, default=8)
parse.add_argument('--test_batch_size', type=int, default=128)
parse.add_argument('--num_workers', type=int, default=32)
parse.add_argument('--eval_fre', type=int, default=1)
parse.add_argument('--msg_fre', type=int, default=10)
parse.add_argument('--save_fre', type=int, default=1)
parse.add_argument('--name', type=str, default='SEResNext50', help='log/model_out/tensorboard log')
parse.add_argument('--data_dir', type=str, default='/media/tiger/zzr/rsna')
parse.add_argument('--log_dir', type=str, default='./logs')
parse.add_argument('--tensorboard_dir', type=str, default='./tensorboard')
parse.add_argument('--model_out_dir', type=str, default='./model_out')
parse.add_argument('--model_out_name', type=str, default='final_model.pth')
parse.add_argument('--seed', type=int, default=5, help='random seed')
parse.add_argument('--predefinedModel', type=str, default='/media/tiger/zzr/rsna_script/model_out/191004-003700_temp/out_1.pth')
return parse.parse_args()
def evalute(net, val_loader, writer, epoch, logger):
logger.info('------------after epo {}, eval...-----------'.format(epoch))
loss = 0
net.eval()
with torch.no_grad():
for img,lb in val_loader:
img, lb = img.cuda(), lb.cuda()
outputs = net(img)
loss += nn.BCELoss()(outputs, lb)
loss /= len(val_loader)
logger.info('loss:{:.4f}/epoch{}'.format(loss, epoch))
writer.add_scalar('loss', loss)
net.train()
def main_worker(args, logger):
try:
writer = SummaryWriter(logdir=args.sub_tensorboard_dir)
train_set = RSDataset(rootpth=args.data_dir, mode='train')
train_loader = DataLoader(train_set,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
pin_memory=True,
num_workers=args.num_workers)
# val_set = RSDataset(rootpth=args.data_dir, mode='train')
# val_loader = DataLoader(val_set,
# batch_size=args.test_batch_size,
# shuffle=False,
# pin_memory=True,
# num_workers=args.num_workers)
net = SEResNext50()
net = net.train()
net = net.cuda()
# net.load_state_dict(torch.load(args.predefinedModel))
criterion = nn.BCELoss().cuda()
# criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.schedule_step, gamma=0.3)
loss_record = []
iter = 0
running_loss = []
st = glob_st = time.time()
total_iter = len(train_loader)*args.epoch
for epoch in range(args.epoch):
# 评估
# evalute(net, val_loader, writer, epoch, logger)
# if epoch != 0 and epoch % args.eval_fre == 0:
# evalute(net, val_loader, writer, epoch, logger)
if epoch != 0 and epoch % args.save_fre == 0:
model_out_name = osp.join(args.sub_model_out_dir,'out_{}.pth'.format(epoch))
# 防止分布式训练保存失败
state_dict = net.modules.state_dict() if hasattr(net, 'module') else net.state_dict()
torch.save(state_dict, model_out_name)
for img, lb in train_loader:
iter += 1
img = img.cuda()
lb = lb.cuda()
optimizer.zero_grad()
outputs = net(img)
loss = criterion(outputs, lb)
loss.backward()
optimizer.step()
running_loss.append(loss.item())
if iter % args.msg_fre == 0:
ed = time.time()
spend = ed-st
global_spend = ed-glob_st
st = ed
eta = int((total_iter-iter)*(global_spend/iter))
eta = str(datetime.timedelta(seconds=eta))
global_spend = str(datetime.timedelta(seconds=(int(global_spend))))
avg_loss = np.mean(running_loss)
loss_record.append(avg_loss)
running_loss = []
lr = optimizer.param_groups[0]['lr']
msg = '. '.join([
'epoch:{epoch}',
'iter/total_iter:{iter}/{total_iter}',
'lr:{lr:.5f}',
'loss:{loss:.4f}',
'spend/global_spend:{spend:.4f}/{global_spend}',
'eta:{eta}'
]).format(
epoch=epoch,
iter=iter,
total_iter=total_iter,
lr=lr,
loss=avg_loss,
spend=spend,
global_spend=global_spend,
eta=eta
)
logger.info(msg)
writer.add_scalar('loss',avg_loss,iter)
writer.add_scalar('lr',lr,iter)
scheduler.step()
# 训练完最后评估一次
# evalute(net, val_loader, writer, args.epoch, logger)
out_name = osp.join(args.sub_model_out_dir,args.model_out_name)
torch.save(net.cpu().state_dict(),out_name)
logger.info('-----------Done!!!----------')
except:
logger.exception('Exception logged')
finally:
writer.close()
if __name__ == '__main__':
args = parse_args()
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# 唯一标识
unique_name = time.strftime('%y%m%d-%H%M%S_') + args.name
args.unique_name = unique_name
# 每次创建作业使用不同的tensorboard目录
args.sub_tensorboard_dir = osp.join(args.tensorboard_dir, args.unique_name)
# 保存模型的目录
args.sub_model_out_dir = osp.join(args.model_out_dir, args.unique_name)
# 创建所有用到的目录
for sub_dir in [args.sub_tensorboard_dir,args.sub_model_out_dir, args.log_dir]:
if not osp.exists(sub_dir):
os.makedirs(sub_dir)
log_file_name = osp.join(args.log_dir,args.unique_name + '.log')
logger = get_logger(log_file_name)
for k, v in args.__dict__.items():
logger.info(k)
logger.info(v)
main_worker(args, logger=logger)
Datasets
Models
