def train(args, model, device, train_loader, optimizer, epoch, joint=True): model.train() focal_loss = FocalLoss(gamma=0.2) for batch_idx, (data, target) in enumerate(train_loader): data = data.to(device) print(data.size()) target = target.to(device) optimizer.zero_grad() output = model(data) if len(output.size()) > 2: loss = 0 for j in range(output.size(1)): if joint: loss += torch.sum( -target * F.log_softmax(output[:, j, :], -1), -1) #target has to be multihot encoding scheme else: loss += F.cross_entropy(output[:, j, :], target) #loss += focal_loss(output[:,j,:],target) loss /= output.size(1) else: loss = F.cross_entropy(output, target) if joint: loss = loss.mean() loss.backward() optimizer.step() if batch_idx % args.log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))
def __init__(self, agent_config): ''' :param agent_config (dict): lr=float,momentum=float,weight_decay=float, schedule=[int], # The last number in the list is the end of epoch model_type=str,model_name=str,out_dim={task:dim},model_weights=str force_single_head=bool print_freq=int gpuid=[int]''' super(Network_, self).__init__() self.log = print self.config = agent_config self.train_loader, self.test_loader = self.config['loaders'][ 'pretrain'] self.model = self.create_model() self.criterion_fn = FocalLoss( ) if self.config['loss'] == 'fl' else CrossEntropyLoss() if agent_config['gpuid'][0] >= 0: self.cuda() self.gpu = True else: self.gpu = False self.exp_name = agent_config['exp_name'] self.init_optimizer() self.n_iter = 0 self.writer = SummaryWriter(log_dir="runs/" + self.exp_name) self.save_after = self.config['save_after']
def __init__(self, task, cuda=True): self.task = 2 self.floattensor = "torch.FloatTensor" if self.task == 2: self.loss_func = FocalLoss(gamma=0.75) #nn.CrossEntropyLoss()# else: self.loss_func = nn.BCEWithLogitsLoss() if cuda: self.loss_func = self.loss_func.cuda() self.floattensor = "torch.cuda.FloatTensor"
def __init__(self, num_labels, dropout_prob, bret_pretrainded_path): """ :param num_labels: :param dropout_prob: :param bret_pretrainded_path: """ # 初始化 super().__init__() # 构建深度学习的网络结构 self.fc = nn.Linear(768, num_labels) # self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=0) self.loss_fn = FocalLoss()
def test(): device = torch.device("cpu") model = LWANet(num_classes=num_classes, pretrained=True) model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(weight_load, map_location=device).items()}) model=model.cuda(device_ids[0]) criterion = FocalLoss(gamma=6) val_file_names = glob.glob('dataset/test/images/*.png') val_dataset = Load_Dataset(val_file_names) val_loader = DataLoader(val_dataset, batch_size=args.batch_size,num_workers=16) val_multi(model, criterion, val_loader, num_classes,batch_size=args.batch_size,device_ids=device_ids)
def __init__(self, net, num_class, classes, size, model_path, multi_task=False, resume=False, always_save=False, use_cuda=True): self.always_save = always_save self.size = size self.classes = classes self.model_path = model_path self.resume = resume self.start_epoch = 0 self.multi_task = multi_task self.best_loss = 1000000 self.net = net self.use_cuda = use_cuda if not net is None and use_cuda: self.net.cuda() self.criterion = FocalLoss(class_num=num_class, alpha=None, gamma=2, size_average=True) if self.resume and self.net is not None: self.start_epoch, self.best_loss = self.resume_model(self.model_path)
def __init__(self, embedding, arch='ExampleNet', dropout_rate=0.2, loss='BCELoss', margin=0, threshold=None, similarity='inner_product', **kwargs): super(ExamplePredictor, self).__init__(**kwargs) self.arch = arch logging.info('building ' + self.arch + '...') self.model = ExampleNet(embedding.size(1), similarity=similarity) self.embedding = torch.nn.Embedding(embedding.size(0), embedding.size(1)) self.embedding.weight = torch.nn.Parameter(embedding) # use cuda self.model = self.model.to(self.device) self.embedding = self.embedding.to(self.device) # make optimizer self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate) self.loss = { 'BCELoss': torch.nn.BCEWithLogitsLoss(), 'FocalLoss': FocalLoss(), #-- 'L1Loss': torch.nn.L1Loss(), 'SmoothL1Loss': torch.nn.SmoothL1Loss(), 'MSELoss': torch.nn.MSELoss(), 'CrossEntropyLoss': torch.nn.CrossEntropyLoss(), 'NLLLoss': torch.nn.NLLLoss(), #'NLLLoss2d': torch.nn.NLLLoss2d(), 'KLDivLoss': torch.nn.KLDivLoss(), 'MarginRankingLoss': torch.nn.MarginRankingLoss(), 'MultiMarginLoss': torch.nn.MultiMarginLoss(), 'MultiLabelMarginLoss': torch.nn.MultiLabelMarginLoss(), 'SoftMarginLoss': torch.nn.SoftMarginLoss(), 'MultiLabelSoftMarginLoss': torch.nn.MultiLabelSoftMarginLoss(), 'CosineEmbeddingLoss': torch.nn.CosineEmbeddingLoss(), 'HingeEmbeddingLoss': torch.nn.HingeEmbeddingLoss(), 'TripleMarginLoss': torch.nn.TripleMarginLoss() }[loss]
def train(): mod = LWANet(num_classes=num_classes, pretrained=True) model = mod.cuda(device_ids[0]) model = nn.DataParallel(model, device_ids=device_ids) batch_size = args.batch_size criterion = FocalLoss(gamma=6) optimizer = optim.Adam(model.parameters(), lr=lra) train_file, val_file = load_filename() liver_dataset = Load_Dataset(train_file) val_dataset = Load_Dataset(val_file) dataloaders = DataLoader(liver_dataset, batch_size=batch_size, shuffle=True, num_workers=12) # drop_last=True val_load = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=12) train_model(model, criterion, optimizer, dataloaders, val_load, num_classes)
def visualize_example(data_loader=None, sample=None, model=None, classifier=None): ''' Input: - data_loader: torch DataLoader, which will generate the plotted image - sample: dictionary specifying {image, anotation} to be plotted - model: torch model to calculate predictions Output: image, groud truth, prediction, superposition of image and prediction ''' if (data_loader is None) and (sample is None): return focallos = FocalLoss(gamma=2) if model is not None: pass #model.eval() #classifier.eval() if data_loader is not None: dataiter = iter(data_loader) image, anotation = dataiter.next() else: image, anotation = sample['image'], sample['anotation'] image = image.to('cuda') anotation = anotation.type('torch.LongTensor') anotation = anotation.to('cuda') if model is not None: with torch.no_grad(): outputs = model(image) #unsqueeze? outputs = classifier(outputs) invTrans = transforms.Compose([ transforms.Normalize(mean=[0., 0., 0.], std=[1 / 0.5, 1 / 0.5, 1 / 0.5]), transforms.Normalize(mean=[-0.5, -0.5, -0.5], std=[1., 1., 1.]), ]) image = invTrans(image[0]) if model is not None: loss = focallos(outputs, anotation.long()) #loss = criterion(outputs, torch.squeeze(anotation,1)) #print(loss.item(), anotation.shape) preds = outputs.argmax(1) #predictions in black & white #F1 score: TP = torch.sum(preds[0] * anotation[0]) FP = torch.sum(preds[0] * ((anotation[0] + 1) % 2)) FN = torch.sum(anotation[0] * ((preds[0] + 1) % 2)) print(TP, FP, FN) p = TP / (TP + FP + 1e-20) r = TP / (TP + FN + 1e-20) F1 = 2 * (p * r) / (p + r + 1e-20) print("precision: ", p) print("recall: ", r) print("F1 score: ", F1) #preds = outputs[:,1,:,:] > 0.3 preds_prob = outputs[:, 1, :, :] #predictions in probability #print(outputs.shape, outputs.argmax(1).shape, "preds shape", preds.shape) preds = preds.to('cpu') preds = preds.type('torch.FloatTensor') preds_PIL = transforms.ToPILImage()(preds[0]) preds_PIL = preds_PIL.convert('RGB') preds_prob = preds_prob.to('cpu') preds_prob = preds_prob.type('torch.FloatTensor') preds_prob_PIL = transforms.ToPILImage()(preds_prob[0]) preds_prob_PIL = preds_prob_PIL.convert('RGB') anotation = anotation.to('cpu') anotation = anotation.type('torch.FloatTensor') #print(anotation[0].dtype) anotation_PIL = transforms.ToPILImage()(anotation[0]) anotation_PIL = anotation_PIL.convert('RGB') image = image.to('cpu') image_PIL = transforms.ToPILImage()(image) image_PIL = image_PIL.convert('RGB') if model is not None: superposition = superposition_anotation( transforms.ToTensor()(image_PIL), transforms.ToTensor()(preds_PIL)) #replace 2 channels to 0s img_grid = torchvision.utils.make_grid([ transforms.ToTensor()(image_PIL), transforms.ToTensor()(anotation_PIL), transforms.ToTensor()(preds_prob_PIL), superposition ]) else: superposition = superposition_anotation( transforms.ToTensor()(image_PIL), transforms.ToTensor()(anotation_PIL)) img_grid = torchvision.utils.make_grid([ transforms.ToTensor()(image_PIL), transforms.ToTensor()(anotation_PIL), superposition ]) # show images matplotlib_imshow(img_grid)
def main(): # 随机种子 np.random.seed(21) torch.manual_seed(21) torch.cuda.manual_seed_all(21) random.seed(21) monitor = pd.DataFrame() # 获取当前文件名,用于创建模型及结果文件的目录 file_name = os.path.basename(__file__).split('.')[0] # 创建保存模型和结果的文件夹 if not os.path.exists('./model/%s' % file_name): os.makedirs('./model/%s' % file_name) if not os.path.exists('./result/%s' % file_name): os.makedirs('./result/%s' % file_name) # 创建日志文件 if not os.path.exists('./result/%s.txt' % file_name): with open('./result/%s.txt' % file_name, 'w') as acc_file: pass with open('./result/%s.txt' % file_name, 'a') as acc_file: acc_file.write('\n%s %s\n' % (time.strftime( "%Y-%m-%d %H:%M:%S", time.localtime(time.time())), file_name)) # 默认使用PIL读图 def default_loader(path): # return Image.open(path) return Image.open(path).convert('RGB') # 训练集图片读取 class TrainDataset(Dataset): def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader): imgs = [] for index, row in label_list.iterrows(): imgs.append((row['img_path'], row['label'])) self.imgs = imgs self.transform = transform self.target_transform = target_transform self.loader = loader def __getitem__(self, index): filename, label = self.imgs[index] img = self.loader(filename) if self.transform is not None: img = self.transform(img) return img, label def __len__(self): return len(self.imgs) # 验证集图片读取 class ValDataset(Dataset): def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader): imgs = [] for index, row in label_list.iterrows(): imgs.append((row['img_path'], row['label'])) self.imgs = imgs self.transform = transform self.target_transform = target_transform self.loader = loader def __getitem__(self, index): filename, label = self.imgs[index] img = self.loader(filename) if self.transform is not None: img = self.transform(img) return img, label def __len__(self): return len(self.imgs) # 测试集图片读取 class TestDataset(Dataset): def __init__(self, label_list, transform=None, target_transform=None, loader=default_loader): imgs = [] for index, row in label_list.iterrows(): imgs.append((row['img_path'])) self.imgs = imgs self.transform = transform self.target_transform = target_transform self.loader = loader def __getitem__(self, index): filename = self.imgs[index] img = self.loader(filename) if self.transform is not None: img = self.transform(img) return img, filename def __len__(self): return len(self.imgs) # 数据增强:在给定角度中随机进行旋转 class FixedRotation(object): def __init__(self, angles): self.angles = angles def __call__(self, img): return fixed_rotate(img, self.angles) def fixed_rotate(img, angles): angles = list(angles) angles_num = len(angles) index = random.randint(0, angles_num - 1) return img.rotate(angles[index]) # 训练函数 def train(train_loader, model, criterion, optimizer, epoch): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() acc = AverageMeter() # switch to train mode model.train() end = time.time() # 从训练集迭代器中获取训练数据 for i, (images, target) in enumerate(train_loader): # 评估图片读取耗时 data_time.update(time.time() - end) # 将图片和标签转化为tensor image_var = torch.tensor(images).cuda(async=True) label = torch.tensor(target).cuda(async=True) # 将图片输入网络,前传,生成预测值 y_pred = model(image_var) # 计算loss loss = criterion(y_pred, label) losses.update(loss.item(), images.size(0)) # 计算top1正确率 prec, PRED_COUNT = accuracy(y_pred.data, target, topk=(1, 1)) acc.update(prec, PRED_COUNT) # 对梯度进行反向传播,使用随机梯度下降更新网络权重 optimizer.zero_grad() loss.backward() optimizer.step() # 评估训练耗时 batch_time.update(time.time() - end) end = time.time() # 打印耗时与结果 if i % print_freq == 0: print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Data {data_time.val:.3f} ({data_time.avg:.3f})\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format( epoch, i, len(train_loader), batch_time=batch_time, data_time=data_time, loss=losses, acc=acc)) return acc.avg, losses.avg # 验证函数 def validate(val_loader, model, criterion): batch_time = AverageMeter() losses = AverageMeter() acc = AverageMeter() # switch to evaluate mode model.eval() end = time.time() for i, (images, labels) in enumerate(val_loader): image_var = torch.tensor(images).cuda(async=True) target = torch.tensor(labels).cuda(async=True) # 图片前传。验证和测试时不需要更新网络权重,所以使用torch.no_grad(),表示不计算梯度 with torch.no_grad(): y_pred = model(image_var) loss = criterion(y_pred, target) # measure accuracy and record loss prec, PRED_COUNT = accuracy(y_pred.data, labels, topk=(1, 1)) losses.update(loss.item(), images.size(0)) acc.update(prec, PRED_COUNT) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % print_freq == 0: print('TrainVal: [{0}/{1}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format( i, len(val_loader), batch_time=batch_time, loss=losses, acc=acc)) print(' * Accuray {acc.avg:.3f}'.format(acc=acc), '(Previous Best Acc: %.3f)' % best_precision, ' * Loss {loss.avg:.3f}'.format(loss=losses), 'Previous Lowest Loss: %.3f)' % lowest_loss) return acc.avg, losses.avg # 测试函数 def test(test_loader, model): csv_map = OrderedDict({'filename': [], 'probability': []}) # switch to evaluate mode model.eval() for i, (images, filepath) in enumerate(tqdm(test_loader)): # bs, ncrops, c, h, w = images.size() filepath = [os.path.basename(i) for i in filepath] image_var = torch.tensor(images, requires_grad=False) # for pytorch 0.4 with torch.no_grad(): y_pred = model(image_var) # 使用softmax函数将图片预测结果转换成类别概率 smax = nn.Softmax(1) smax_out = smax(y_pred) # 保存图片名称与预测概率 csv_map['filename'].extend(filepath) for output in smax_out: prob = ';'.join([str(i) for i in output.data.tolist()]) csv_map['probability'].append(prob) result = pd.DataFrame(csv_map) result['probability'] = result['probability'].map( lambda x: [float(i) for i in x.split(';')]) # 转换成提交样例中的格式 sub_filename, sub_label = [], [] for index, row in result.iterrows(): sub_filename.append(row['filename']) pred_label = np.argmax(row['probability']) if pred_label == 0: sub_label.append('norm') else: sub_label.append('defect%d' % pred_label) # 生成结果文件,保存在result文件夹中,可用于直接提交 submission = pd.DataFrame({ 'filename': sub_filename, 'label': sub_label }) submission.to_csv('./result/%s/submission.csv' % file_name, header=None, index=False) print('test done!') return # 保存最新模型以及最优模型 def save_checkpoint(state, is_best, is_lowest_loss, filename='./model/%s/checkpoint.pth.tar' % file_name): torch.save(state, filename) if is_best: shutil.copyfile(filename, './model/%s/model_best.pth.tar' % file_name) if is_lowest_loss: shutil.copyfile(filename, './model/%s/lowest_loss.pth.tar' % file_name) # 用于计算精度和时间的变化 class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self): self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count # 学习率衰减:lr = lr / lr_decay def adjust_learning_rate(): nonlocal lr lr = lr / lr_decay return optim.Adam(model.parameters(), lr, weight_decay=weight_decay, amsgrad=True) # 计算top K准确率 def accuracy(y_pred, y_actual, topk=(1, )): """Computes the precision@k for the specified values of k""" final_acc = 0 maxk = max(topk) # for prob_threshold in np.arange(0, 1, 0.01): PRED_COUNT = y_actual.size(0) PRED_CORRECT_COUNT = 0 prob, pred = y_pred.topk(maxk, 1, True, True) # prob = np.where(prob > prob_threshold, prob, 0) for j in range(pred.size(0)): if int(y_actual[j]) == int(pred[j]): PRED_CORRECT_COUNT += 1 if PRED_COUNT == 0: final_acc = 0 else: final_acc = PRED_CORRECT_COUNT / PRED_COUNT return final_acc * 100, PRED_COUNT # 程序主体 # 设定GPU ID os.environ["CUDA_VISIBLE_DEVICES"] = '2,3' # 小数据集上,batch size不易过大。如出现out of memory,应调小batch size batch_size = 6 # 进程数量,最好不要超过电脑最大进程数,尽量能被batch size整除。windows下报错可以改为workers=0 workers = 12 # epoch数量,分stage进行,跑完一个stage后降低学习率进入下一个stage stage_epochs = [20, 10, 10, 10, 10] # 初始学习率 lr = 1e-4 # 学习率衰减系数 (new_lr = lr / lr_decay) lr_decay = 5 # 正则化系数 weight_decay = 1e-4 # 原始正则化系数 1e-4 # 参数初始化 stage = 0 start_epoch = 0 total_epochs = sum(stage_epochs) best_precision = 0 lowest_loss = 100 # 设定打印频率,即多少step打印一次,用于观察loss和acc的实时变化 # 打印结果中,括号前面为实时loss和acc,括号内部为epoch内平均loss和acc print_freq = 10 # 验证集比例 val_ratio = 0.12 # 是否只验证,不训练 evaluate = False # 是否从断点继续跑 resume = False # 创建inception_v4模型 # model = model_v4.v4(num_classes=12) model = make_model( 'pnasnet5large', pretrained=True, num_classes=12, dropout_p=0.5, ) model = torch.nn.DataParallel(model).cuda() # optionally resume from a checkpoint if resume: checkpoint_path = './model/%s/checkpoint.pth.tar' % file_name if os.path.isfile(checkpoint_path): print("=> loading checkpoint '{}'".format(checkpoint_path)) checkpoint = torch.load(checkpoint_path) start_epoch = checkpoint['epoch'] + 1 best_precision = checkpoint['best_precision'] lowest_loss = checkpoint['lowest_loss'] stage = checkpoint['stage'] lr = checkpoint['lr'] model.load_state_dict(checkpoint['state_dict']) # 如果中断点恰好为转换stage的点,需要特殊处理 if start_epoch in np.cumsum(stage_epochs)[:-1]: stage += 1 optimizer = adjust_learning_rate() model.load_state_dict( torch.load('./model/%s/model_best.pth.tar' % file_name)['state_dict']) print("=> loaded checkpoint (epoch {})".format( checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(resume)) # 读取训练图片列表 all_data = pd.read_csv('data/label.csv') # 分离训练集和测试集,stratify参数用于分层抽样 train_data_list, val_data_list = train_test_split( all_data, test_size=val_ratio, random_state=21, shuffle=True, stratify=all_data['label']) # 读取测试图片列表 test_data_list = pd.read_csv('data/test.csv') # 图片归一化,由于采用ImageNet预训练网络,因此这里直接采用ImageNet网络的参数 normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 训练集图片变换,输入网络的尺寸为384*384 train_data = TrainDataset( train_data_list, transform=transforms.Compose([ transforms.Resize((400, 400)), transforms.ColorJitter(0.15, 0.15, 0.15, 0.075), transforms.RandomHorizontalFlip(), transforms.RandomGrayscale(), # transforms.RandomRotation(20), FixedRotation([0, 30, 45, 60, 90, 135, 180, 270]), # 添加了45、135 两个角度 ljc transforms.RandomCrop(384), transforms.ToTensor(), normalize, ])) # 验证集图片变换 val_data = ValDataset(val_data_list, transform=transforms.Compose([ transforms.Resize((400, 400)), transforms.CenterCrop(384), transforms.ToTensor(), normalize, ])) # 测试集图片变换 test_data = TestDataset(test_data_list, transform=transforms.Compose([ transforms.Resize((400, 400)), transforms.CenterCrop(384), transforms.ToTensor(), normalize, ])) # 生成图片迭代器 train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=workers) val_loader = DataLoader(val_data, batch_size=batch_size * 2, shuffle=False, pin_memory=False, num_workers=workers) test_loader = DataLoader(test_data, batch_size=batch_size * 2, shuffle=False, pin_memory=False, num_workers=workers) # 对占比较少的类别进行加权 # weight = torch.Tensor([1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5]) # 使用交叉熵损失函数 # weight added by ljc # weight = torch.Tensor([1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 10]) # criterion = nn.CrossEntropyLoss(weight=weight).cuda() # weight criterion = FocalLoss(class_num=12, gamma=2) # changed by ljc # criterion = nn.CrossEntropyLoss().cuda() # weight # 优化器,使用带amsgrad的Adam optimizer = optim.Adam(model.parameters(), lr, weight_decay=weight_decay, amsgrad=True) if evaluate: validate(val_loader, model, criterion) else: # 开始训练 for epoch in range(start_epoch, total_epochs): # train for one epoch train_acc, train_loss = train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set precision, avg_loss = validate(val_loader, model, criterion) monitor = pd.concat([ monitor, pd.DataFrame({ 'train_acc': [train_acc], 'train_loss': [train_loss], 'val_acc': [precision], 'val_acc': [avg_loss] }) ]) # 在日志文件中记录每个epoch的精度和loss with open('./result/%s.txt' % file_name, 'a') as acc_file: acc_file.write('Epoch: %2d, Precision: %.8f, Loss: %.8f\n' % (epoch, precision, avg_loss)) # 记录最高精度与最低loss,保存最新模型与最佳模型 is_best = precision > best_precision is_lowest_loss = avg_loss < lowest_loss best_precision = max(precision, best_precision) lowest_loss = min(avg_loss, lowest_loss) state = { 'epoch': epoch, 'state_dict': model.state_dict(), 'best_precision': best_precision, 'lowest_loss': lowest_loss, 'stage': stage, 'lr': lr, } save_checkpoint(state, is_best, is_lowest_loss) # 判断是否进行下一个stage if (epoch + 1) in np.cumsum(stage_epochs)[:-1]: stage += 1 optimizer = adjust_learning_rate() model.load_state_dict( torch.load('./model/%s/model_best.pth.tar' % file_name)['state_dict']) print('Step into next stage') with open('./result/%s.txt' % file_name, 'a') as acc_file: acc_file.write( '---------------Step into next stage----------------\n' ) # 记录线下最佳分数 with open('./result/%s.txt' % file_name, 'a') as acc_file: acc_file.write('* best acc: %.8f %s\n' % (best_precision, os.path.basename(__file__))) with open('./result/best_acc.txt', 'a') as acc_file: acc_file.write( '%s * best acc: %.8f %s\n' % (time.strftime("%Y-%m-%d %H:%M:%S", time.localtime( time.time())), best_precision, os.path.basename(__file__))) # 读取最佳模型,预测测试集,并生成可直接提交的结果文件 best_model = torch.load('./model/%s/model_best.pth.tar' % file_name) model.load_state_dict(best_model['state_dict']) test(test_loader=test_loader, model=model) # 释放GPU缓存 torch.cuda.empty_cache()
def _train_or_test(model, dataloader, config: Settings, optimizer=None, use_l1_mask=True, log_writer: SummaryWriter = None, step: int = 0, weighting_attention=False, is_valid=False): ''' model: the multi-gpu model dataloader: optimizer: if None, will be test evaluation ''' is_train = optimizer is not None n_examples = 0 n_correct = 0 n_batches = 0 total_cross_entropy = 0 total_cluster_cost = 0 # separation cost is meaningful only for class_specific total_separation_cost = 0 total_avg_separation_cost = 0 total_loss = 0 conf_matrix = np.zeros((2, 2), dtype='int32') preds = [] targets = [] if config.loss_function == 'cross_entropy': loss_fn = torch.nn.CrossEntropyLoss() elif config.loss_function == 'focal': loss_fn = FocalLoss(alpha=0.5, gamma=2) else: raise NotImplementedError('unknown loss function: ' + config.loss_function) with tqdm(total=len(dataloader.dataset), unit='bag') as pbar: for i, (image, label) in enumerate(dataloader): input = image.cuda() # if param all_labels=True in dataloader, set label to positive if at least one in the list if len(label) > 1: label = label.max().unsqueeze(0) target = label.cuda() # torch.enable_grad() has no effect outside of no_grad() grad_req = torch.enable_grad() if is_train else torch.no_grad() with grad_req: output, min_distances, attention, _ = model.forward_(input) cross_entropy = loss_fn(output, target) if config.mil_pooling == 'loss_attention': instance_labels = target * torch.ones( input.size(0), dtype=torch.long, device=input.device) loss_2 = WeightCrossEntropyLoss()(model.out_c, instance_labels, model.A) cross_entropy += 2.0 * loss_2 if config.class_specific: max_dist = (model.prototype_shape[1] * model.prototype_shape[2] * model.prototype_shape[3]) # prototypes_of_correct_class is a tensor of shape batch_size * num_prototypes attention_detached = attention.detach().cpu() weight = np.interp( attention_detached, (attention_detached.min(), attention_detached.max()), (0.001, 1)) if weighting_attention: tensor_weight = torch.tensor(weight).cuda() else: tensor_weight = torch.tensor(1).cuda() # calculate cluster cost prototypes_of_correct_class = torch.t( model.prototype_class_identity[:, label]).cuda() inverted_distances, _ = torch.max( (max_dist - (min_distances * tensor_weight.T)) * prototypes_of_correct_class, dim=1) cluster_cost = torch.mean(max_dist - inverted_distances) # calculate separation cost prototypes_of_wrong_class = 1 - prototypes_of_correct_class inverted_distances_to_nontarget_prototypes, _ = \ torch.max((max_dist - (min_distances * tensor_weight.T)) * prototypes_of_wrong_class, dim=1) separation_cost = torch.mean( max_dist - inverted_distances_to_nontarget_prototypes) # calculate avg cluster cost avg_separation_cost = \ torch.sum((min_distances * tensor_weight.T) * prototypes_of_wrong_class, dim=1) / torch.sum( prototypes_of_wrong_class, dim=1) avg_separation_cost = torch.mean(avg_separation_cost) if use_l1_mask: l1_mask = 1 - torch.t( model.prototype_class_identity).cuda() l1 = (model.last_layer.weight * l1_mask).norm(p=1) else: l1 = model.last_layer.weight.norm(p=1) else: min_distance, _ = torch.min(min_distances, dim=1) cluster_cost = torch.mean(min_distance) l1 = model.last_layer.weight.norm(p=1) # evaluation statistics _, predicted = torch.max(output.data, 1) n_examples += target.size(0) n_correct += (predicted == target).sum().item() pred_s = func.softmax(output, dim=-1) preds.append(pred_s.data.cpu().numpy()) targets.append(target.cpu().numpy()) conf_matrix += confusion_matrix(target.cpu().numpy(), predicted.cpu().numpy(), labels=[0, 1]) n_batches += 1 total_cross_entropy += cross_entropy.item() total_cluster_cost += cluster_cost.item() total_separation_cost += separation_cost.item() total_avg_separation_cost += avg_separation_cost.item() # compute gradient and do SGD step if config.class_specific: loss = (config.coef_crs_ent * cross_entropy + config.coef_clst * cluster_cost + config.coef_sep * separation_cost + config.coef_l1 * l1) else: loss = (config.coef_crs_ent * cross_entropy + config.coef_clst * cluster_cost + config.coef_l1 * l1) total_loss += loss.item() if is_train: optimizer.zero_grad() loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5) # gradient clipping optimizer.step() del input del target del output del predicted del min_distances pbar.update(1) total_cross_entropy /= n_batches total_cluster_cost /= n_batches total_separation_cost /= n_batches total_loss /= n_batches preds = np.concatenate(preds) targets = np.concatenate(targets) auc = roc_auc_score(targets, preds[..., 1]) pred_y = preds.argmax(1) precision = precision_score(targets, pred_y, zero_division=0) recall = recall_score(targets, pred_y, zero_division=0) f1 = f1_score(targets, pred_y, zero_division=0) fpr, tpr, threshold = roc_curve(targets, preds[..., 1]) print('\t\taccuracy:', n_correct / n_examples) print('\t\tauc:', auc) print('\t\ttotal_loss:', total_loss) if is_train: suffix = '/train' else: if is_valid: suffix = '/valid' else: suffix = '/test' if log_writer: log_writer.add_scalar('total_loss' + suffix, total_loss, global_step=step) log_writer.add_scalar('cross_entropy' + suffix, total_cross_entropy, global_step=step) log_writer.add_scalar('cluster_cost' + suffix, total_cluster_cost, global_step=step) if config.class_specific: log_writer.add_scalar('separation_cost' + suffix, total_separation_cost, global_step=step) log_writer.add_scalar('avg_separation_cost' + suffix, total_avg_separation_cost / n_batches, global_step=step) log_writer.add_scalar('accuracy' + suffix, n_correct / n_examples, global_step=step) log_writer.add_scalar('auc' + suffix, auc, global_step=step) log_writer.add_scalar('precision' + suffix, precision, global_step=step) log_writer.add_scalar('recall' + suffix, recall, global_step=step) log_writer.add_scalar('f-score' + suffix, f1, global_step=step) log_writer.add_scalar('l1' + suffix, model.last_layer.weight.norm(p=1).item(), global_step=step) conf_plot = ConfusionMatrixDisplay(confusion_matrix=conf_matrix).plot( cmap='Blues', values_format='d') log_writer.add_figure('confusion_matrix' + suffix, conf_plot.figure_, global_step=step, close=True) plt.figure() plt.title('Receiver Operating Characteristic') plt.plot(fpr, tpr, 'b', label='AUC = %0.2f' % auc) plt.legend(loc='lower right') plt.plot([0, 1], [0, 1], 'r--') plt.xlim([0, 1]) plt.ylim([0, 1]) plt.ylabel('True Positive Rate') plt.xlabel('False Positive Rate') log_writer.add_figure('roc' + suffix, plt.gcf(), global_step=step, close=True) p = model.prototype_vectors.view(model.num_prototypes, -1).cpu() with torch.no_grad(): p_avg_pair_dist = torch.mean(list_of_distances(p, p)) if log_writer: log_writer.add_scalar('p_avg_pair_dist' + suffix, p_avg_pair_dist, global_step=step) return auc
opt.scales = [opt.initial_scale] for i in range(1, opt.n_scales): opt.scales.append(opt.scales[-1] * opt.scale_step) opt.arch = '{}-{}'.format(opt.model, opt.model_depth) opt.mean = get_mean(opt.norm_value, dataset=opt.mean_dataset) opt.std = get_std(opt.norm_value) print(opt) with open(os.path.join(opt.result_path, 'opts.json'), 'w') as opt_file: json.dump(vars(opt), opt_file) torch.manual_seed(opt.manual_seed) model, parameters = generate_model(opt) print(model) criterion = FocalLoss(num_class=opt.n_classes) #nn.CrossEntropyLoss() if not opt.no_cuda: criterion = criterion.cuda() if opt.no_mean_norm and not opt.std_norm: norm_method = Normalize([0, 0, 0], [1, 1, 1]) elif not opt.std_norm: norm_method = Normalize(opt.mean, [1, 1, 1]) else: norm_method = Normalize(opt.mean, opt.std) if not opt.no_train: assert opt.train_crop in ['random', 'corner', 'center'] if opt.train_crop == 'random': crop_method = MultiScaleRandomCrop(opt.scales, opt.sample_size) elif opt.train_crop == 'corner':
if ntokens > 500000: # One Billion # This produces fairly even matrix mults for the buckets: # 0: 11723136, 1: 10854630, 2: 11270961, 3: 11219422 splits = [4200, 35000, 180000] elif ntokens > 75000: # WikiText-103 splits = [2800, 20000, 76000] print('Using', splits) if args.loss == 'splitcrossentropy': criterion = SplitCrossEntropyLoss(args.emsize, splits=splits, verbose=False) fn_exclude_keys += ['gamma'] elif args.loss == 'focal': criterion = FocalLoss(args.gamma) fn_exclude_keys += ['emsize'] ### if args.cuda: model = model.cuda() criterion = criterion.cuda() ### params = list(model.parameters()) + list(criterion.parameters()) total_params = sum(x.size()[0] * x.size()[1] if len(x.size()) > 1 else x.size()[0] for x in params if x.size()) print('Args:', args) print('Model total parameters:', total_params) ## get filename fn = 'save/' + get_name(args.__dict__,
def main(): torch.backends.cudnn.benchmark = True global best_loss, args, use_gpu global step, test_loader args = parser.parse_args() use_gpu = torch.cuda.is_available() train_imgdir = '/home/zhaoyang/data/voc0712/train' test_imgdir = '/home/zhaoyang/data/voc0712/test' train_annotation_file = '/home/zhaoyang/data/voc0712/annotation/train_annotation.txt' test_annotaiion_file = '/home/zhaoyang/data/voc0712/annotation/test_annotation.txt' normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_set = VocDataset(train_imgdir, train_annotation_file, input_size=[600, 600], transform=transforms.Compose( [transforms.ToTensor(), normalizer])) train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True, num_workers=8, collate_fn=train_set.collate_fn) test_set = VocDataset(test_imgdir, test_annotaiion_file, input_size=[600, 600], transform=transforms.Compose( [transforms.ToTensor(), normalizer])) test_loader = torch.utils.data.DataLoader(test_set, batch_size=4, shuffle=True, num_workers=8, collate_fn=test_set.collate_fn) model = retinanet101(pretrained=True, num_classes=args.num_classes) model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3]) optimizer = optim.SGD(model.parameters(), 0.001, momentum=args.momentum, weight_decay=args.weight_decay) #lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 40], gamma=0.1) #criterion = FocalLoss(args.alpha * torch.ones(args.num_classes, 1), args.gamma) criterion = FocalLoss() if args.resume: if os.path.isfile(args.resume): print('Loading model from {}'.format(args.resume)) checkpoint = torch.load(args.resume) model = model.load_state_dict(checkpoint['model']) optimizer = optimizer.load_state_dict((checkpoint['optimizer'])) args.start_step = checkpoint['step'] best_loss = checkpoint['best_loss'] print('Loaded model from {} (epoch {})'.format( args.resume, args.start_epoch)) else: print('No checkpoint founded in {}'.format(args.resume)) if use_gpu: model.cuda() criterion.cuda() step = args.start_step while step <= args.steps: #lr_scheduler.step() train_model(model, train_loader, optimizer, criterion)
def main(): global args, best_prec1 best_prec1 = 1e6 args = parser.parse_args() args.original_lr = 1e-6 args.lr = 1e-6 args.momentum = 0.95 args.decay = 5 * 1e-4 args.start_epoch = 0 args.epochs = 5000 args.steps = [-1, 1, 100, 150] args.scales = [1, 1, 1, 1] args.workers = 4 args.seed = time.time() args.print_freq = 30 wandb.config.update(args) wandb.run.name = f"Default_{wandb.run.name}" if ( args.task == wandb.run.name) else f"{args.task}_{wandb.run.name}" conf = configparser.ConfigParser() conf.read(args.config) print(conf) TRAIN_DIR = conf.get("COVNet_CSR_Focal_CBAM_HIGH", "train") VALID_DIR = conf.get("COVNet_CSR_Focal_CBAM_HIGH", "valid") TEST_DIR = conf.get("COVNet_CSR_Focal_CBAM_HIGH", "test") LOG_DIR = conf.get("COVNet_CSR_Focal_CBAM_HIGH", "log") create_dir_not_exist(LOG_DIR) train_list = [ os.path.join(TRAIN_DIR, item) for item in os.listdir(TRAIN_DIR) ] val_list = [ os.path.join(VALID_DIR, item) for item in os.listdir(VALID_DIR) ] os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu torch.cuda.manual_seed(args.seed) model = COVNet_CSR_CBAM(3) model = model.cuda() criterion = FocalLoss(num_class=3, size_average=False).cuda() optimizer = torch.optim.Adam(model.parameters(), args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay) model = DataParallel_withLoss(model, criterion) for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) train(train_list, model, criterion, optimizer, epoch) prec1 = validate(val_list, model, criterion, epoch) with open(os.path.join(LOG_DIR, args.task + ".txt"), "a") as f: f.write("epoch " + str(epoch) + " BCELoss: " + str(float(prec1))) f.write("\n") wandb.save(os.path.join(LOG_DIR, args.task + ".txt")) is_best = prec1 < best_prec1 best_prec1 = min(prec1, best_prec1) print(' * best BCELoss {BCELoss:.3f} '.format(BCELoss=best_prec1)) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.pre, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best, args.task, epoch=epoch, path=os.path.join(LOG_DIR, args.task))
def main(args): best_acc = -1 logger = bit_common.setup_logger(args) cp, cn = smooth_BCE(eps=0.1) # Lets cuDNN benchmark conv implementations and choose the fastest. # Only good if sizes stay the same within the main loop! torch.backends.cudnn.benchmark = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"Going to train on {device}") classes = 5 train_set, valid_set, train_loader, valid_loader = mktrainval(args, logger) print(len(train_loader)) logger.info(f"Loading model from {args.model}.npz") model = models.KNOWN_MODELS[args.model](head_size=classes, zero_head=True) model.load_from(np.load(f"{args.model}.npz")) logger.info("Moving model onto all GPUs") model = torch.nn.DataParallel(model) # Optionally resume from a checkpoint. # Load it to CPU first as we'll move the model to GPU later. # This way, we save a little bit of GPU memory when loading. start_epoch = 0 # Note: no weight-decay! optim = torch.optim.SGD(model.parameters(), lr=0.003, momentum=0.9) # Resume fine-tuning if we find a saved model. savename = pjoin(args.logdir, args.name, "bit.pth.tar") try: logger.info(f"Model will be saved in '{savename}'") checkpoint = torch.load(savename, map_location="cpu") logger.info(f"Found saved model to resume from at '{savename}'") start_epoch = checkpoint["epoch"] model.load_state_dict(checkpoint["model"]) optim.load_state_dict(checkpoint["optim"]) logger.info(f"Resumed at epoch {start_epoch}") except FileNotFoundError: logger.info("Fine-tuning from BiT") model = model.to(device) optim.zero_grad() model.train() #mixup = bit_hyperrule.get_mixup(len(train_set)) mixup = -1 from focalloss import FocalLoss cri = FocalLoss(gamma=1.0) #cri = torch.nn.CrossEntropyLoss().to(device) logger.info("Starting training!") chrono = lb.Chrono() accum_steps = 0 mixup_l = np.random.beta(mixup, mixup) if mixup > 0 else 1 end = time.time() epoches = 20 scheduler = torch.optim.lr_scheduler.OneCycleLR(optim, max_lr=0.01, steps_per_epoch=1, epochs=epoches) with lb.Uninterrupt() as u: for epoch in range(start_epoch, epoches): pbar = enumerate(train_loader) pbar = tqdm.tqdm(pbar, total=len(train_loader)) scheduler.step() all_top1, all_top5 = [], [] for param_group in optim.param_groups: lr = param_group["lr"] #for x, y in recycle(train_loader): for batch_id, (x, y) in pbar: #for batch_id, (x, y) in enumerate(train_loader): # measure data loading time, which is spent in the `for` statement. chrono._done("load", time.time() - end) if u.interrupted: break # Schedule sending to GPU(s) x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) # Update learning-rate, including stop training if over. #lr = bit_hyperrule.get_lr(step, len(train_set), args.base_lr) #if lr is None: # break if mixup > 0.0: x, y_a, y_b = mixup_data(x, y, mixup_l) # compute output with chrono.measure("fprop"): logits = model(x) top1, top5 = topk(logits, y, ks=(1, 5)) all_top1.extend(top1.cpu()) all_top5.extend(top5.cpu()) if mixup > 0.0: c = mixup_criterion(cri, logits, y_a, y_b, mixup_l) else: c = cri(logits, y) train_loss = c.item() train_acc = np.mean(all_top1)*100.0 # Accumulate grads with chrono.measure("grads"): (c / args.batch_split).backward() accum_steps += 1 accstep = f"({accum_steps}/{args.batch_split})" if args.batch_split > 1 else "" s = f"epoch={epoch} batch {batch_id}{accstep}: loss={train_loss:.5f} train_acc={train_acc:.2f} lr={lr:.1e}" #s = f"epoch={epoch} batch {batch_id}{accstep}: loss={c.item():.5f} lr={lr:.1e}" pbar.set_description(s) #logger.info(f"[batch {batch_id}{accstep}]: loss={c_num:.5f} (lr={lr:.1e})") # pylint: disable=logging-format-interpolation logger.flush() # Update params with chrono.measure("update"): optim.step() optim.zero_grad() # Sample new mixup ratio for next batch mixup_l = np.random.beta(mixup, mixup) if mixup > 0 else 1 # Run evaluation and save the model. val_loss, val_acc = run_eval(model, valid_loader, device, chrono, logger, epoch) best = val_acc > best_acc if best: best_acc = val_acc torch.save({ "epoch": epoch, "val_loss": val_loss, "val_acc": val_acc, "train_acc": train_acc, "model": model.state_dict(), "optim" : optim.state_dict(), }, savename) end = time.time() logger.info(f"Timings:\n{chrono}")
def unet_train(): batch_size = 1 num_epochs = [5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000] num_workers = 2 lr = 0.0001 losslist = ['dice'] # ['focal', 'bce', 'dice', 'lovasz'] optimlist = ['adam'] # ['adam', 'sgd'] iflog = True SC_root_dir = '../dataset-EdmSealedCrack-512' train_files, val_files, test_files = myutils.organize_SC_files(SC_root_dir) train_RC_dataset = DatasetRealCrack('../dataset-EdmCrack600-512/A/train', transform=transform) train_SC_dataset = DatasetSealedCrack(files=train_files, root_dir=SC_root_dir, transform=data_Train_transforms) val_RC_dataset = DatasetRealCrack('../dataset-EdmCrack600-512/A/val', transform=transform) val_SC_dataset = DatasetSealedCrack(files=val_files, root_dir=SC_root_dir, transform=data_Test_transforms) train_loader = torch.utils.data.DataLoader(ConcatDataset( train_RC_dataset, train_SC_dataset), batch_size=2, shuffle=True, num_workers=2) criterion = nn.BCELoss() focallos = FocalLoss(gamma=2) doubleFocalloss = focalloss.FocalLoss_2_datasets(gamma=2) epoidx = -1 for los in losslist: for opt in optimlist: start = time.time() print(los, opt) torch.manual_seed(77) torch.cuda.manual_seed(77) ################# #unet = Unet_SpatialPyramidPooling(3).cuda() ################# unet = Unet(3).cuda() SC_classifier = classifier(64, 2).cuda() RC_classifier = classifier(64, 2).cuda() ################## #unet = smp.Unet('resnet34', encoder_weights='imagenet').cuda() #unet.segmentation_head = torch.nn.Sequential().cuda() #SC_classifier = classifier(16, 2).cuda() #RC_classifier = classifier(16, 2).cuda() #UNCOMMENT TO KEEP TRAINING THE BEST MODEL prev_epoch = 0 # if loading model 58, change to prev_epoch = 58. When saving the model, it is going to be named as 59, 60, 61... #unet.load_state_dict(torch.load('trained_models/unet_adam_dice_58.pkl')) #SC_classifier.load_state_dict(torch.load('trained_models/SC_classifier_adam_dice_58.pkl')) #RC_classifier.load_state_dict(torch.load('trained_models/RC_classifier_adam_dice_58.pkl')) history = [] if 'adam' in opt: optimizer = torch.optim.Adam(unet.parameters(), lr=lr) elif 'sgd' in opt: optimizer = torch.optim.SGD(unet.parameters(), lr=10 * lr, momentum=0.9) logging.basicConfig(filename='./logs/logger_unet.log', level=logging.INFO) total_step = len(train_loader) epoidx += 1 for epoch in range(num_epochs[epoidx]): totalloss = 0 for i, (realCrack_batch, sealedCrack_batch) in enumerate(train_loader): SC_images = sealedCrack_batch[0].cuda() SC_masks = sealedCrack_batch[1].cuda() RC_images = realCrack_batch[0].cuda() RC_masks = realCrack_batch[1].cuda() SC_encoder = unet(SC_images) RC_encoder = unet(RC_images) ############# SC_outputs = SC_classifier(SC_encoder) RC_outputs = RC_classifier(RC_encoder) ############# #Deep lab v3 #SC_outputs = SC_classifier(SC_encoder['out']) #RC_outputs = RC_classifier(RC_encoder['out']) ############## if 'bce' in los: masks = onehot(masks) loss = criterion(outputs, masks) elif 'dice' in los: branch_RC = {'outputs': RC_outputs, 'masks': RC_masks} branch_SC = {'outputs': SC_outputs, 'masks': SC_masks} loss = dice_loss_2_datasets(branch_RC, branch_SC) #masks = onehot(masks) #loss = dice_loss(outputs, masks) elif 'lovasz' in los: masks = onehot(masks) loss = L.lovasz_hinge(outputs, masks) elif 'focal' in los: #loss = focallos(outputs, masks.long()) branch_RC = { 'outputs': RC_outputs, 'masks': RC_masks.long() } branch_SC = { 'outputs': SC_outputs, 'masks': SC_masks.long() } loss = doubleFocalloss(branch_RC, branch_SC) totalloss += loss * RC_images.size(0) #*2? #print(RC_images.size(0)) optimizer.zero_grad() loss.backward() optimizer.step() if i % 10 == 0: print(epoch, i) print("total loss: ", totalloss) if i % 1000 == 0: print("Epoch:%d; Iteration:%d; Loss:%f" % (epoch, i, loss)) if i + 1 == total_step: # and epoch%1==0: #and val_miou>0.85: torch.save( unet.state_dict(), './trained_models/unet_' + opt + '_' + los + '_' + str(epoch + 1 + prev_epoch) + '.pkl') torch.save( RC_classifier.state_dict(), './trained_models/RC_classifier_' + opt + '_' + los + '_' + str(epoch + 1 + prev_epoch) + '.pkl') torch.save( SC_classifier.state_dict(), './trained_models/SC_classifier_' + opt + '_' + los + '_' + str(epoch + 1 + prev_epoch) + '.pkl') history_np = np.array(history) np.save('./logs/unet_' + opt + '_' + los + '.npy', history_np) end = time.time() print((end - start) / 60)
def main_worker(index, opt): random.seed(opt.manual_seed) np.random.seed(opt.manual_seed) torch.manual_seed(opt.manual_seed) if index >= 0 and opt.device.type == 'cuda': opt.device = torch.device(f'cuda:{index}') if opt.distributed: opt.dist_rank = opt.dist_rank * opt.ngpus_per_node + index dist.init_process_group(backend='nccl', init_method=opt.dist_url, world_size=opt.world_size, rank=opt.dist_rank) opt.batch_size = int(opt.batch_size / opt.ngpus_per_node) opt.n_threads = int( (opt.n_threads + opt.ngpus_per_node - 1) / opt.ngpus_per_node) opt.is_master_node = not opt.distributed or opt.dist_rank == 0 model = generate_model(opt) if opt.batchnorm_sync: assert opt.distributed, 'SyncBatchNorm only supports DistributedDataParallel.' model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) if opt.pretrain_path: model = load_pretrained_model(model, opt.pretrain_path, opt.model, opt.n_finetune_classes) if opt.resume_path is not None: model = resume_model(opt.resume_path, opt.arch, model) model = make_data_parallel(model, opt.distributed, opt.device) if opt.pretrain_path: parameters = get_fine_tuning_parameters(model, opt.ft_begin_module) else: parameters = model.parameters() if opt.is_master_node: print(model) #criterion = CrossEntropyLoss().to(opt.device) # ADDED for 231n criterion = FocalLoss().to(opt.device) if not opt.no_train: (train_loader, train_sampler, train_logger, train_batch_logger, optimizer, scheduler) = get_train_utils(opt, parameters) if opt.resume_path is not None: opt.begin_epoch, optimizer, scheduler = resume_train_utils( opt.resume_path, opt.begin_epoch, optimizer, scheduler) if opt.overwrite_milestones: scheduler.milestones = opt.multistep_milestones if not opt.no_val: val_loader, val_logger = get_val_utils(opt) if opt.tensorboard and opt.is_master_node: from torch.utils.tensorboard import SummaryWriter if opt.begin_epoch == 1: tb_writer = SummaryWriter(log_dir=opt.result_path) else: tb_writer = SummaryWriter(log_dir=opt.result_path, purge_step=opt.begin_epoch) else: tb_writer = None conf_mtx_dict = {} # ADDED for CS231n prev_val_loss = None for i in range(opt.begin_epoch, opt.n_epochs + 1): if not opt.no_train: if opt.distributed: train_sampler.set_epoch(i) current_lr = get_lr(optimizer) train_epoch(i, train_loader, model, criterion, optimizer, opt.device, current_lr, train_logger, train_batch_logger, tb_writer, opt.distributed) if i % opt.checkpoint == 0 and opt.is_master_node: save_file_path = opt.result_path / 'save_{}.pth'.format(i) save_checkpoint(save_file_path, i, opt.arch, model, optimizer, scheduler) if not opt.no_val: prev_val_loss = val_epoch(i, val_loader, model, criterion, opt.device, val_logger, tb_writer, opt.distributed, conf_mtx_dict) # ADDED for CS231n # ADDED for 231n - uncomment if using cross entropy loss #if not opt.no_train and opt.lr_scheduler == 'multistep': # scheduler.step() #elif not opt.no_train and opt.lr_scheduler == 'plateau': # scheduler.step(prev_val_loss) if opt.inference: inference_loader, inference_class_names = get_inference_utils(opt) inference_result_path = opt.result_path / '{}.json'.format( opt.inference_subset) inference.inference(inference_loader, model, inference_result_path, inference_class_names, opt.inference_no_average, opt.output_topk) # ADDED for CS231n conf_mtx_file = csv.writer(open("conf_mtxs.csv", "w+")) for key, val in conf_mtx_dict.items(): conf_mtx_file.writerow([key, val])
def main(): torch.backends.cudnn.benchmark = True global best_loss, args, use_gpu args = parser.parse_args() use_gpu = torch.cuda.is_available() train_imgdir = '/home/zhaoyang/data/voc0712/train' test_imgdir = '/home/zhaoyang/data/voc0712/test' train_annotation_file = '/home/zhaoyang/data/voc0712/annotation/train_annotation.txt' test_annotaiion_file = '/home/zhaoyang/data/voc0712/annotation/test_annotation.txt' normalizer = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_set = VocDataset(train_imgdir, train_annotation_file, input_size=[600, 600], transform=transforms.Compose( [transforms.ToTensor(), normalizer])) train_loader = torch.utils.data.DataLoader(train_set, batch_size=8, shuffle=True, num_workers=8, collate_fn=train_set.collate_fn) test_set = VocDataset(test_imgdir, test_annotaiion_file, input_size=[600, 600], transform=transforms.Compose( [transforms.ToTensor(), normalizer])) test_loader = torch.utils.data.DataLoader(test_set, batch_size=4, shuffle=True, num_workers=8, collate_fn=test_set.collate_fn) model = retinanet101(pretrained=True, num_classes=args.num_classes) # model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count())) optimizer = optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[80, 120], gamma=0.1) #criterion = FocalLoss(args.alpha * torch.ones(args.num_classes, 1), args.gamma) criterion = FocalLoss() if args.resume: if os.path.isfile(args.resume): print('Loading model from {}'.format(args.resume)) checkpoint = torch.load(args.resume) model = model.load_state_dict(checkpoint['model']) optimizer = optimizer.load_state_dict((checkpoint['optimizer'])) args.start_epoch = checkpoint['epoch'] best_loss = checkpoint['best_loss'] print('Loaded model from {} (epoch {})'.format( args.resume, args.start_epoch)) else: print('No checkpoint founded in {}'.format(args.resume)) if use_gpu: model.cuda() criterion.cuda() for epoch in xrange(args.start_epoch, args.epochs): #test_loss = test_model(model, test_loader, criterion) #print("test loss", test_loss) lr_scheduler.step() train_model(model, train_loader, optimizer, criterion, epoch) if (epoch + 1) % args.test_freq == 0: test_loss = test_model(model, test_loader, criterion) if test_loss <= best_loss: best_loss = test_loss save_checkpoint( { 'epoch': epoch + 1, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_loss': best_loss }, epoch + 1, True) print('best test loss: {}'.format(best_loss)) if (epoch + 1) % args.save_freq == 0: save_checkpoint( { 'epoch': epoch + 1, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_loss': best_loss }, epoch + 1, False)
def run_experiment(_exp_name, _epochs, _train_manifest, _test_manifest, _labels, _use_mfcc_in, _use_ivectors_in, _use_embeddings_in, _use_transcripts_out, _use_accents_out, _batch_size, _num_workers, _mfcc_size, _ivector_size, _embedding_size, _rnn_type, _rnn_hidden_size, _nb_head_layers, _nb_speech_layers, _nb_accents_layers, _bidirectional, _losses_mix, _learning_rate, _lm_path, _decoder_alpha, _decoder_beta, _decoder_cutoff_top_n, _decoder_beam_width, _cuda, _tensorboard_path, _saved_models_path, _bottleneck_size, _accent_loss): print(f'\n##### Running experiment {_exp_name} #####') # Tools to log values results_dict = {} results_dict['train_loss'] = [] results_dict['train_loss_text'] = [] results_dict['train_loss_accent'] = [] results_dict['test_loss'] = [] results_dict['test_loss_text'] = [] results_dict['test_loss_accent'] = [] results_dict['test_wer'] = [] results_dict['test_accent_acc'] = [] tb_path = Path(_tensorboard_path) / _exp_name makedirs(tb_path, exist_ok=True) tb_writer = SummaryWriter(tb_path) ### DATA LOADING # Training set train_dataset = MultiDataset(_train_manifest, _labels, use_mfcc_in=_use_mfcc_in, use_ivectors_in=_use_ivectors_in, use_embeddings_in=_use_embeddings_in, embedding_size=_embedding_size, use_transcripts_out=_use_transcripts_out, use_accents_out=_use_accents_out) train_loader = MultiDataLoader(train_dataset, batch_size=_batch_size, shuffle=True, num_workers=_num_workers) # Testing set test_dataset = MultiDataset(_test_manifest, _labels, use_mfcc_in=_use_mfcc_in, use_ivectors_in=_use_ivectors_in, use_embeddings_in=_use_embeddings_in, embedding_size=_embedding_size, use_transcripts_out=_use_transcripts_out, use_accents_out=_use_accents_out) test_loader = MultiDataLoader(test_dataset, batch_size=_batch_size, shuffle=True, num_workers=_num_workers) ### CREATE MODEL model = MultiTask(use_mfcc_in=_use_mfcc_in, use_ivectors_in=_use_ivectors_in, use_embeddings_in=_use_embeddings_in, use_transcripts_out=_use_transcripts_out, use_accents_out=_use_accents_out, mfcc_size=_mfcc_size, ivector_size=_ivector_size, embedding_size=_embedding_size, rnn_type=_rnn_type, labels=_labels, accents_dict=train_dataset.accent_dict, rnn_hidden_size=_rnn_hidden_size, nb_head_layers=_nb_head_layers, nb_speech_layers=_nb_speech_layers, nb_accents_layers=_nb_accents_layers, bidirectional=_bidirectional, bottleneck_size=_bottleneck_size, DEBUG=False) if _cuda: model = model.cuda() print(model, '\n') print('Model parameters counts:', MultiTask.get_param_size(model), '\n') ### OPTIMIZER, CRITERION, DECODER # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=_learning_rate) # Criterion if _use_accents_out: if _accent_loss == 'focal': AccLoss = FocalLoss() elif _accent_loss == 'CE': AccLoss = nn.CrossEntropyLoss() else: raise ValueError( f'Loss {_accent_loss} for accent_loss is unknown. Please use either "focal" or "CE".' ) if not _use_transcripts_out: # only accent classification criterion = AccLoss elif not _use_accents_out: # only text recognition criterion = nn.CTCLoss() else: # both tasks criterion = (nn.CTCLoss(), FocalLoss()) # Decoder if _use_transcripts_out: decoder = BeamCTCDecoder(_labels, lm_path=_lm_path, alpha=_decoder_alpha, beta=_decoder_beta, cutoff_top_n=_decoder_cutoff_top_n, cutoff_prob=_decoder_cutoff_top_n, beam_width=_decoder_beam_width, num_processes=_num_workers) target_decoder = GreedyDecoder(_labels) else: decoder, target_decoder = None, None ### EPOCHS best_wer = math.inf best_acc = 0 for epoch in range(1, _epochs + 1): ### TRAIN print(f'Epoch {epoch} training: {exp_name}') train_results = train(model, train_loader, criterion, optimizer, losses_mix=_losses_mix) train_loss, train_loss_text, train_loss_accent = train_results results_dict['train_loss'].append(train_loss) results_dict['train_loss_text'].append(train_loss_text) results_dict['train_loss_accent'].append(train_loss_accent) print(f'Epoch {epoch} training loss: {train_loss}') ### TEST print(f'Epoch {epoch} testing') test_results = test(model, test_loader, criterion, decoder, target_decoder, losses_mix=_losses_mix) test_loss, test_loss_text, test_loss_accent, test_wer, test_accent_acc = test_results results_dict['test_loss'].append(test_loss) results_dict['test_loss_text'].append(test_loss_text) results_dict['test_loss_accent'].append(test_loss_accent) results_dict['test_wer'].append(test_wer) results_dict['test_accent_acc'].append(test_accent_acc) print(f'Epoch {epoch} testing loss: {test_loss}') # Add values to tensorboard for key, results in results_dict.items(): tb_writer.add_scalar(key, results[-1], epoch) #Save model if it is best save_new = False if _use_transcripts_out: if test_wer < best_wer: save_new = True best_wer = test_wer else: if test_accent_acc > best_acc: save_new = True best_acc = test_accent_acc if save_new: MultiTask.serialize( model, Path(_saved_models_path) / _exp_name, save=True, exp_name=_exp_name, optimizer=optimizer, epoch=epoch, train_losses=results_dict['train_loss'], test_losses=results_dict['test_loss'], text_train_losses=results_dict['train_loss_text'], text_test_losses=results_dict['test_loss_text'], text_wers=results_dict['test_wer'], accent_train_losses=results_dict['train_loss_accent'], accent_test_losses=results_dict['test_loss_accent'], accent_accuracies=results_dict['test_accent_acc']) del model gc.collect() torch.cuda.empty_cache()
def forward( self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): r""" labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels - 1]``. """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.albert( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict, ) sequence_output = outputs[2] layers = len(sequence_output) batchsize, length, hidden_size = sequence_output[0].size( 0), sequence_output[0].size(1), sequence_output[0].size(2) sequence_output = torch.cat(sequence_output).view( layers, batchsize, length, hidden_size) sequence_output = sequence_output.transpose(0, 1).transpose( 1, 2).contiguous() sequence_output = self.attn(sequence_output) if self.quick_return: return sequence_output sequence_output = self.dropout(sequence_output) logits = self.classifier(sequence_output) loss = None if labels is not None: if self.lossfct == 'diceloss': loss_fct = MultiDiceLoss() if attention_mask is not None: active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels) active_labels = labels.view(-1) active_labels = F.one_hot(active_labels, self.num_labels) mask = attention_mask.view(-1, 1) mask = mask.repeat(1, self.num_labels) loss = loss_fct(active_logits, active_labels, mask) #print(loss) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) elif self.lossfct == 'focalloss': loss_fct = FocalLoss() # Only keep active parts of the loss if attention_mask is not None: active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels) active_labels = torch.where( active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)) loss = loss_fct(active_logits, active_labels) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) else: loss_fct = CrossEntropyLoss(reduction=self.CEL_type) # Only keep active parts of the loss if attention_mask is not None: active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels) active_labels = torch.where( active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)) loss = loss_fct(active_logits, active_labels) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) if not return_dict: output = (logits, ) + outputs[2:] return ((loss, ) + output) if loss is not None else output return TokenClassifierOutput( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )
def forward( self, input_ids=None, attention_mask=None, mems=None, perm_mask=None, target_mapping=None, token_type_ids=None, input_mask=None, head_mask=None, inputs_embeds=None, labels=None, #use_mems=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs, ): r""" labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): Labels for computing the multiple choice classification loss. Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension of the input tensors. (see `input_ids` above) """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.transformer( input_ids, attention_mask=attention_mask, mems=mems, perm_mask=perm_mask, target_mapping=target_mapping, token_type_ids=token_type_ids, input_mask=input_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, #use_mems=use_mems, output_attentions=output_attentions, output_hidden_states=True, return_dict=return_dict, ) sequence_output = outputs.hidden_states # tuple( [batch, len, hiddenstate],) layers = len(sequence_output) batchsize, length, hidden_size = sequence_output[0].size( 0), sequence_output[0].size(1), sequence_output[0].size(2) '''print(layers) print(batchsize) print(length) print(hidden_size)''' # print(sequence_output.size()) sequence_output = torch.cat(sequence_output).view( layers, batchsize, length, hidden_size) # tensor.size([layers, batch, len, hiddenstate]) # print(sequence_output.size()) sequence_output = sequence_output.transpose(0, 1).transpose( 1, 2).contiguous() sequence_output = self.attn(sequence_output) if self.quick_return: return sequence_output logits = self.classifier(sequence_output) loss = None if labels is not None: if self.lossfct == 'diceloss': loss_fct = MultiDiceLoss() if attention_mask is not None: '''print(attention_mask) print(attention_mask.shape) #torch.Size([4, 80]) batch,len''' active_loss = attention_mask.view(-1) == 1 '''print(active_loss) print(active_loss.shape)#torch.Size([320]) 4*80 print(logits) print(logits.shape) #torch.Size([4, 80, 6])''' active_logits = logits.view(-1, self.num_labels) '''print(active_logits) print(active_logits.shape)#torch.Size([320, 6]) 4*80*6''' #active_logits = torch.masked_select(active_logits, (active_loss == 1)) active_labels = labels.view(-1) #->torch.Size([320]) active_labels = F.one_hot(active_labels, self.num_labels) '''print(labels) print(labels.shape)#torch.Size([4, 80]) print(active_labels) print(active_labels.shape)#torch.Size([320,6]) print(active_logits)''' mask = attention_mask.view(-1, 1) mask = mask.repeat(1, self.num_labels) '''print(mask) print(mask.shape)#torch.Size([320, 6])''' loss = loss_fct(active_logits, active_labels, mask) #print(loss) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) elif self.lossfct == 'focalloss': loss_fct = FocalLoss() # 'sum' # Only keep active parts of the loss if attention_mask is not None: active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels) active_labels = torch.where( active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)) loss = loss_fct(active_logits, active_labels) # 320*6, 320 else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) else: loss_fct = CrossEntropyLoss(reduction=self.CEL_type) #'sum' # Only keep active parts of the loss if attention_mask is not None: active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels) active_labels = torch.where( active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)) loss = loss_fct(active_logits, active_labels) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) if not return_dict: output = (logits, ) + outputs[1:] return ((loss, ) + output) if loss is not None else output return XLNetForTokenClassificationOutput( loss=loss, logits=logits, mems=outputs.mems, hidden_states=outputs.hidden_states, attentions=outputs.attentions, )
def train_src_threemodal(model1, model2, model3, train_loader1, train_loader2, train_loader3, val_loader): global lr global best_prec1 lr = params.base_lr # model1 = construct_resnet18(model1, params) # model2 = construct_resnet18(model2, params) # model3 = construct_resnet18(model3, params) model1 = construct_resnet34(model1, params) model2 = construct_resnet34(model2, params) model3 = construct_resnet34(model3, params) model1.train() model2.train() model3.train() optimizer1 = torch.optim.Adam(list(model1.parameters()), lr=params.base_lr, betas=(0.9, 0.99)) optimizer2 = torch.optim.Adam(list(model2.parameters()), lr=params.base_lr, betas=(0.9, 0.99)) optimizer3 = torch.optim.Adam(list(model3.parameters()), lr=params.base_lr, betas=(0.9, 0.99)) # criterion = nn.CrossEntropyLoss().cuda() focalloss = FocalLoss(gamma=2) for epoch in range(params.start_epoch, params.start_epoch + params.num_epochs): adjust_learning_rate(optimizer1, epoch, params.base_lr) adjust_learning_rate(optimizer2, epoch, params.base_lr) adjust_learning_rate(optimizer3, epoch, params.base_lr) # train for one epoch # train_batch(train_loader, model, criterion, optimizer, epoch) for step, (images, labels) in enumerate(train_loader1): # make images and labels variable images = make_variable(images) labels = make_variable(labels.squeeze_()) # zero gradients for optimizer optimizer1.zero_grad() # compute loss for critic preds = model1(images) loss = focalloss(preds, labels) # optimize source classifier loss.backward() optimizer1.step() # print step info if ((step + 1) % params.log_step_pre == 0): print("Color Epoch [{}/{}] Step [{}/{}]: loss={}".format( epoch + 1, params.num_epochs, step + 1, len(train_loader1), loss.item())) for step, (images, labels) in enumerate(train_loader2): # make images and labels variable images = make_variable(images) labels = make_variable(labels.squeeze_()) # zero gradients for optimizer optimizer2.zero_grad() # compute loss for critic preds = model2(images) loss = focalloss(preds, labels) # optimize source classifier loss.backward() optimizer2.step() # print step info if ((step + 1) % params.log_step_pre == 0): print("Depth Epoch [{}/{}] Step [{}/{}]: loss={}".format( epoch + 1, params.num_epochs, step + 1, len(train_loader2), loss.item())) for step, (images, labels) in enumerate(train_loader3): # make images and labels variable images = make_variable(images) labels = make_variable(labels.squeeze_()) # zero gradients for optimizer optimizer3.zero_grad() # compute loss for critic preds = model3(images) loss = focalloss(preds, labels) # optimize source classifier loss.backward() optimizer3.step() # print step info if ((step + 1) % params.log_step_pre == 0): print("Ir Epoch [{}/{}] Step [{}/{}]: loss={}".format( epoch + 1, params.num_epochs, step + 1, len(train_loader3), loss.item())) if ((epoch + 1) % params.eval_step_pre == 0): eval_acc(model1, model2, model3, val_loader) # save model parameters if ((epoch + 1) % params.save_step_pre == 0): save_model(model1, "MultiNet-color-{}.pt".format(epoch + 1)) save_model(model2, "MultiNet-depth-{}.pt".format(epoch + 1)) save_model(model3, "MultiNet-ir-{}.pt".format(epoch + 1)) # # save final model save_model(model1, "MultiNet-color-final.pt") save_model(model2, "MultiNet-depth-final.pt") save_model(model3, "MultiNet-ir-final.pt") return model1, model2, model3
def train_net(args): cropsize = [cfgs.crop_height, cfgs.crop_width] # dataset_train = CityScapes(cfgs.data_dir, cropsize=cropsize, mode='train') dataset_train = ContextVoc(cfgs.train_file, cropsize=cropsize, mode='train') dataloader_train = DataLoader(dataset_train, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, drop_last=True) # dataset_val = CityScapes(cfgs.data_dir, mode='val') dataset_val = ContextVoc(cfgs.val_file, cropsize=cropsize, mode='train') dataloader_val = DataLoader(dataset_val, batch_size=1, shuffle=True, num_workers=args.num_workers, drop_last=True) # build net os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda if torch.cuda.is_available() and args.use_gpu: device = torch.device('cuda') else: device = torch.device('cpu') # model = BiSeNet(args.num_classes, args.context_path) net = DeeplabV3plus(cfgs).to(device) # net = SCAR(load_weights=True).to(device) if args.pretrained_model_path is not None: print('load model from %s ...' % args.pretrained_model_path) state_dict = torch.load(args.pretrained_model_path, map_location=device) state_dict = renamedict(state_dict) net.load_state_dict(state_dict, strict=False) # net.load_state_dict(torch.load(args.pretrained_model_path)) print('Done!') if args.mulgpu: net = torch.nn.DataParallel(net) net.train() # build optimizer if args.optimizer == 'rmsprop': optimizer = torch.optim.RMSprop(net.parameters(), args.learning_rate) elif args.optimizer == 'sgd': optimizer = torch.optim.SGD(net.parameters(), args.learning_rate, momentum=0.9, weight_decay=1e-4) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(net.parameters(), args.learning_rate) else: print('not supported optimizer \n') optimizer = None #build loss if args.losstype == 'dice': criterion = DiceLoss() elif args.losstype == 'crossentropy': criterion = torch.nn.CrossEntropyLoss() elif args.losstype == 'ohem': score_thres = 0.7 n_min = args.batch_size * cfgs.crop_height * cfgs.crop_width // 16 criterion = OhemCELoss(thresh=score_thres, n_min=n_min) elif args.losstype == 'focal': # criterion = SoftmaxFocalLoss() criterion = FocalLoss() elif args.losstype == 'multi': criterion = Multiloss(4) return net, optimizer, criterion, dataloader_train, dataloader_val