typing_extensions出现如下问题
NameError: name ‘doc’ is not defined
大概为typing_extensions的343行那里的doc出了问题,
将所有的代码替换为如下,问题就解决了。`
import os
import time
import argparse
import numpy as np
import torch.multiprocessing as mp
import torch.nn.functional as F
import torch
# from networks.resnet50_3d_gcn_x5 import RESNET50_3D_GCN_X5
# from torchvision.models.resnet import resnet50
from msaspp_dataloader import msasppDataLoader, colorize_mask
from torch.utils.tensorboard import SummaryWriter
from sklearn.metrics import confusion_matrix
from torchvision import transforms
from torch.backends import cudnn
from denseASPP import DenseASPP
from cfgs import DenseASPP121
from tqdm import tqdm
parser = argparse.ArgumentParser(description='Multi-scale ASPP training')
parser.add_argument('--mode', type=str, help='training and validation mode', default='train')
parser.add_argument('--model_name', type=str, help='model name to be trained', default='denseaspp-v3')
# Dataset
parser.add_argument('--data_path', type=str, help='training data path', default=os.getcwd())
parser.add_argument('--input_height', type=int, help='input height', default=512)
parser.add_argument('--input_width', type=int, help='input width', default=512)
# Training
parser.add_argument('--num_seed', type=int, help='random seed number', default=1)
parser.add_argument('--batch_size', type=int, help='train batch size', default=8)
parser.add_argument('--num_epochs', type=int, help='number of epochs', default=80)
parser.add_argument('--learning_rate', type=float, help='initial learning rate', default=3e-4)
parser.add_argument('--weight_decay', type=float, help='weight decay factor for optimization', default=1e-5)
parser.add_argument('--retrain', type=bool, help='If used with checkpoint_path, will restart training from step zero', default=False)
parser.add_argument('--do_eval', type=bool, help='Mod to evaluating the training model', default=False)
# Preprocessing
parser.add_argument('--random_rotate', type=bool, help='if set, will perform random rotation for augmentation', default=False)
parser.add_argument('--degree', type=float, help='random rotation maximum degree', default=2.5)
# Log and save
parser.add_argument('--checkpoint_path', type=str, help='path to a specific checkpoint to load', default='')
parser.add_argument('--log_directory', type=str, help='directory to save checkpoints and summaries', default=os.path.join(os.getcwd(), 'log'))
parser.add_argument('--log_freq', type=int, help='Logging frequency in global steps', default=250)
parser.add_argument('--save_freq', type=int, help='Checkpoint saving frequency in global steps', default=500)
# Multi-gpu training
parser.add_argument('--gpu', type=int, help='GPU id to use', default=0)
parser.add_argument('--rank', type=int, help='node rank(tensor dimension)for distributed training', default=0)
parser.add_argument('--dist_url', type=str, help='url used to set up distributed training', default='file:///c:/MultiGPU.txt')
parser.add_argument('--dist_backend', type=str, help='distributed backend', default='gloo')
parser.add_argument('--num_threads', type=int, help='number of threads to use for data loading', default=5)
parser.add_argument('--world_size', type=int, help='number of nodes for distributed training', default=1)
parser.add_argument('--multiprocessing_distributed', help='Use multi-processing distributed training to launch '
'N process per node, which has N GPUs. '
'This is the fastest way to use PyTorch for either single node or '
'multi node data parallel training', default=False)
args = parser.parse_args()
def check_folder(path):
if not os.path.exists(path):
os.makedirs(path)
def evaluate_model(val_dataloader, model, criterion):
eval_iou_sum_score = 0
eval_loss_sum = 0
for sampled_eval in tqdm(val_dataloader.data):
with torch.no_grad():
eval_image = sampled_eval['image'].cuda(args.gpu, non_blocking=True)
eval_gt = sampled_eval['gt'].cuda(args.gpu, non_blocking=True)
eval_output = model(eval_image)
model.eval()
eval_loss = criterion(eval_output, eval_gt)
eval_output = F.softmax(eval_output, dim=1)
eval_output = torch.argmax(eval_output, dim=1)
eval_output = eval_output.contiguous().view(-1)
eval_gt = eval_gt.contiguous().view(-1)
iou_per_class = []
for num_class in range(len(val_dataloader.class_names)):
true_class = (eval_output == num_class)
true_label = (eval_gt == num_class)
if true_label.long().sum().item() == 0:
iou_per_class.append(np.nan)
else:
intersect = torch.logical_and(true_class, true_label).sum().float().item()
union = torch.logical_or(true_class, true_label).sum().float().item()
iou = (intersect + 1e-10) / (union + 1e-10)
iou_per_class.append(iou)
eval_iou_sum_score += np.nanmean(iou_per_class)
eval_loss_sum += eval_loss
return eval_loss_sum, eval_iou_sum_score
def main():
model_filename = args.model_name + '.py'
command = 'mkdir ' + os.path.join(args.log_directory, args.model_name)
os.system(command)
command = 'mkdir ' + os.path.join(args.log_directory, args.model_name, 'model')
os.system(command)
command = 'mkdir ' + os.path.join(args.log_directory, args.model_name, 'eval_model')
os.system(command)
if args.checkpoint_path == '':
aux_out_path = os.path.join(args.log_directory, args.model_name)
command = 'cp denseASPP.py ' + aux_out_path
os.system(command)
command = 'cp msaspp_main.py ' + aux_out_path
os.system(command)
command = 'cp msaspp_dataloader.py ' + aux_out_path
os.system(command)
else:
loaded_model_dir = os.path.dirname(args.checkpoint_path)
loaded_model_name = os.path.basename(loaded_model_dir)
loaded_model_filename = loaded_model_name + '.py'
model_out_path = os.path.join(args.checkpoint_path, args.model_name, model_filename)
torch.cuda.empty_cache()
args.distributed = args.
world_size > 1 or args.multiprocessing_distributed
ngpus_per_node = torch.cuda.device_count()
if args.multiprocessing_distributed:
args.world_size = ngpus_per_node * args.world_size
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
else:
main_worker(ngpus_per_node, args)
def main_worker(ngpus_per_node, args):
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
model = DenseASPP(args, model_cfg=DenseASPP121.Model_CFG)
model.train()
num_params = sum([np.prod(p.size()) for p in model.parameters()])
print("Total number of parameters: {}".format(num_params))
num_params_update = sum([np.prod(p.shape) for p in model.parameters() if p.requires_grad])
print("Total number of learning parameters: {}".format(num_params_update))
if args.distributed:
a=1
else:
model = torch.nn.DataParallel(model)
model.cuda()
if args.distributed:
print("Model Initialized on GPU: {}".format(args.gpu))
else:
print("Model Initialized")
global_step = 0
optimizer = torch.optim.Adam(model.parameters(), weight_decay=args.weight_decay, lr=args.learning_rate)
if args.checkpoint_path != '':
if os.path.isfile(args.checkpoint_path):
print("Loading checkpoint '{}'".format(args.checkpoint_path))
if args.gpu in None:
checkpoint = torch.load(args.checkpoint_path)
else:
loc = 'cuda: {}'.format(args.gpu)
checkpoint = torch.load(args.checkpoint_path, map_location=loc)
global_step = checkpoint['global_step']
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("Loaded checkpoint '{}' (global_step {})".format(args.checkpoint_path, checkpoint['global_step']))
else:
print("No checkpoint found at '{}'".format(args.checkpoint_path))
if args.retrain:
global_step = 0
cudnn.benchmark = True
dataloader = msasppDataLoader(args, mode='train')
val_dataloader = msasppDataLoader(args, mode='val')
# Logging
writer = SummaryWriter(os.path.join(args.log_directory, args.model_name, 'summaries'), flush_secs=30)
criterion = torch.nn.CrossEntropyLoss(ignore_index=dataloader.ignore_label).cuda()
start_time = time.time()
duration = 0
steps_per_epoch = len(dataloader.data)
num_total_steps = args.num_epochs * steps_per_epoch
epoch = global_step // steps_per_epoch
mIoU_list = []
train_loss_list = []
while epoch < args.num_epochs:
if args.distributed:
dataloader.train_sampler.set_epoch(epoch)
for step, sample_batched in enumerate(dataloader.data):
optimizer.zero_grad()
befor_op_time = time.time()
sample_image = sample_batched['image'].cuda(args.gpu, non_blocking=True)
sample_gt = sample_batched['gt'].cuda(args.gpu, non_blocking=True)
output = model(sample_image)
loss = criterion(output, sample_gt)
loss.backward()
for param_group in optimizer.param_groups:
current_lr = args.learning_rate * ((1 - epoch / args.num_epochs) ** 0.9)
param_group['lr'] = current_lr
optimizer.step()
if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):
print('[epoch][s/s_per_e/global_step]: [{}/{}][{}/{}/{}], lr: {:.12f}, train loss: {:.12f}'.format(epoch+1, args.num_epochs, step+1, steps_per_epoch, global_step+1, current_lr, loss))
if np.isnan(loss.cpu().item()):
print('NaN in loss occurred. Aborting training.')
return -1
train_loss_list.append(loss)
duration += time.time() - befor_op_time
if global_step and global_step % args.log_freq == 0:
examples_per_sec = args.batch_size / duration * args.log_freq
duration = 0
time_sofar = (time.time() - start_time) / 3600
training_time_left = (num_total_steps / global_step - 1.0) * time_sofar
# if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):
# print("{}".format(args.model_name))
print_string = 'GPU: {} | examples/s: {:4.2f} | Average train loss: {:.5f} | time elapsed: {:.2f}h | time left: {:.2f}h'
print(print_string.format(args.gpu, examples_per_sec, sum(train_loss_list)/len(train_loss_list), time_sofar, training_time_left))
if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):
writer.add_scalar('loss', loss, global_step=global_step)
writer.add_scalar('learning_rate', current_lr, global_step=global_step)
output = output.data.cpu().numpy()
output = output.argmax(axis=1)
for i in range(args.batch_size):
color_gt = colorize_mask(output[i, :]).convert('RGB')
color_gt = transforms.ToTensor()(color_gt)
writer.add_image('segmentation_gt/image/{}'.format(i), torch.unsqueeze(sample_gt, 1)[i, :].data, global_step=global_step)
writer.add_image('segmentation_est/image/{}'.format(i), color_gt.data, global_step=global_step)
writer.flush()
checkpoint = {'global_step': global_step,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()}
torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'model', 'model-{:07d}.pth'.format(global_step)))
if global_step and global_step % args.save_freq == 0:
if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):
eval_loss_sum, eval_iou_sum_score = evaluate_model(val_dataloader, model, criterion)
mIoU = eval_iou_sum_score / len(val_dataloader.data)
mIoU_list.append(mIoU)
idx_max = mIoU_list.index(max(mIoU_list))
if idx_max == 0 and idx_max == (global_step // args.save_freq - 1):
print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f}'
print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), eval_loss_sum/len(val_dataloader.data), mIoU_list[idx_max]))
checkpoint = {'global_step': global_step,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'eval_mIoU': mIoU_list[idx_max]}
torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'eval_model', 'model-{:07d}_mIoU-{:.3f}.pth'.format(global_step, mIoU_list[idx_max])))
elif idx_max == 0 and idx_max != (global_step // args.save_freq - 1):
pass
elif idx_max !=0 and idx_max == (global_step // args.save_freq - 1):
print_string = 'GPU: {} | Epoch: {}/{} | Average train loss: {:.5f} | Average validation loss: {:.5f} | evaluation mIoU: {:.5f}'
print(print_string.format(args.gpu, epoch, args.num_epochs, sum(train_loss_list)/len(train_loss_list), eval_loss_sum/len(val_dataloader.data), max(mIoU_list)))
checkpoint = {'global_step': global_step,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'eval_mIoU': max(mIoU_list)}
torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'eval_model', 'model-{:07d}_mIoU-{:.3f}.pth'.format(global_step, max(mIoU_list))))
model.train()
global_step += 1
epoch += 1
checkpoint = {'global_step': global_step,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()}
torch.save(checkpoint, os.path.join(args.log_directory, args.model_name, 'model', 'model-{:07d}.pth'.format(global_step)))
if __name__ == '__main__':
main()
在这里插入代码片`