train_partseg
日志代码没有解释，已在train_classification中，且该代码是在GPU下运行
点击查看代码
"""
Author: Benny
Date: Nov 2019
"""
import argparse
import os
import torch
import datetime
import logging
import sys
import importlib
import shutil
import provider
import numpy as np
from pathlib import Path
from tqdm import tqdm
from data_utils.ShapeNetDataLoader import PartNormalDataset
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))
seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37],
'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49],
'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}
seg_label_to_cat = {}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat #seg_label_to_cat也是一个字典，但此时键是标签，值是类名，形如{0:Airplane, 1:Airplane, ...49:Table}
def inplace_relu(m): #如train_classification中的作用是节约内存
classname = m.__class__.__name__
if classname.find('ReLU') != -1:
m.inplace=True
def to_categorical(y, num_classes):
""" 1-hot encodes a tensor """
new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
if (y.is_cuda):
return new_y.cuda()
return new_y
def parse_args():
parser = argparse.ArgumentParser('Model')
parser.add_argument('--model', type=str, default='pointnet_part_seg', help='model name')
parser.add_argument('--batch_size', type=int, default=16, help='batch Size during training')
parser.add_argument('--epoch', default=251, type=int, help='epoch to run')
parser.add_argument('--learning_rate', default=0.001, type=float, help='initial learning rate')
parser.add_argument('--gpu', type=str, default='0', help='specify GPU devices')
parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD')
parser.add_argument('--log_dir', type=str, default=None, help='log path')
parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay')
parser.add_argument('--npoint', type=int, default=2048, help='point Number')
parser.add_argument('--normal', action='store_true', default=False, help='use normals')
parser.add_argument('--step_size', type=int, default=20, help='decay step for lr decay')
parser.add_argument('--lr_decay', type=float, default=0.5, help='decay rate for lr decay')
return parser.parse_args()
def main(args):
def log_string(str):
logger.info(str)
print(str)
'''HYPER PARAMETER'''
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
'''CREATE DIR'''
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
exp_dir = Path('./log/')
exp_dir.mkdir(exist_ok=True)
exp_dir = exp_dir.joinpath('part_seg')
exp_dir.mkdir(exist_ok=True)
if args.log_dir is None:
exp_dir = exp_dir.joinpath(timestr)
else:
exp_dir = exp_dir.joinpath(args.log_dir)
exp_dir.mkdir(exist_ok=True)
checkpoints_dir = exp_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = exp_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
'''LOG'''
args = parse_args()
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)
root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/'
TRAIN_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='trainval', normal_channel=args.normal)
trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=10, drop_last=True)
TEST_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='test', normal_channel=args.normal)
testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=10)
log_string("The number of training data is: %d" % len(TRAIN_DATASET))
log_string("The number of test data is: %d" % len(TEST_DATASET))
num_classes = 16
num_part = 50
'''MODEL LOADING'''
MODEL = importlib.import_module(args.model)
shutil.copy('models/%s.py' % args.model, str(exp_dir))
shutil.copy('models/pointnet2_utils.py', str(exp_dir))
classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda()
criterion = MODEL.get_loss().cuda()
classifier.apply(inplace_relu)
def weights_init(m): #权值初始化
classname = m.__class__.__name__ #conv1和Linear都是类，且这两个module都有weight属性
if classname.find('Conv1d') != -1: #如果类名当中有Conv1或者Linear的话,这里自己将conv2改成了conv1,应该是代码作者自己的失误
torch.nn.init.xavier_normal_(m.weight.data) #xavier初始化方法中服从正态分布
torch.nn.init.constant_(m.bias.data, 0.0) #...(Tensor,val)-使用val的值来填充输入的Tensor
elif classname.find('Linear') != -1:
torch.nn.init.xavier_normal_(m.weight.data)
torch.nn.init.constant_(m.bias.data, 0.0)
try:
checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth')
start_epoch = checkpoint['epoch']
classifier.load_state_dict(checkpoint['model_state_dict'])
log_string('Use pretrain model')
except:
log_string('No existing model, starting training from scratch...')
start_epoch = 0
classifier = classifier.apply(weights_init) #classifier.apply(weights_init)对classifier里的每一个module都施加weights_init的作用
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(
classifier.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=args.decay_rate
)
else:
optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9)
def bn_momentum_adjust(m, momentum):
if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d):
m.momentum = momentum #BatchNorm1d中的momentum：更新全局均值running_mean和方差running_var时使用该值进行平滑
LEARNING_RATE_CLIP = 1e-5
MOMENTUM_ORIGINAL = 0.1
MOMENTUM_DECCAY = 0.5
MOMENTUM_DECCAY_STEP = args.step_size
best_acc = 0
global_epoch = 0
best_class_avg_iou = 0
best_inctance_avg_iou = 0
for epoch in range(start_epoch, args.epoch):
mean_correct = []
'''Setting hyperparameters'''
log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch))
'''Adjust learning rate and BN momentum'''
lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP) #该句代码保证虽然学习率下降，但不会低于LEARNING_RATE_CLIP
log_string('Learning rate:%f' % lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP)) #//-表示整除取因子
if momentum < 0.01: #保证momentum值不低于0.01
momentum = 0.01
print('BN momentum updated to: %f' % momentum)
classifier = classifier.apply(lambda x: bn_momentum_adjust(x, momentum))
classifier = classifier.train()
'''learning one epoch'''
for i, (points, label, target) in tqdm(enumerate(trainDataLoader), total=len(trainDataLoader), smoothing=0.9):
optimizer.zero_grad()
points = points.data.numpy()
points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
points = torch.Tensor(points)
points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, trans_feat = classifier(points, to_categorical(label, num_classes)) #label对应于cls，num_classes对应于seg
seg_pred = seg_pred.contiguous().view(-1, num_part) #contiguous()-类似深拷贝,此时的seg_pred的维度是[B*N,num_part]
target = target.view(-1, 1)[:, 0] #此时target的维度是[B*N,1]
pred_choice = seg_pred.data.max(1)[1] #此时pred_choice的维度是[B*N,1]
correct = pred_choice.eq(target.data).cpu().sum()
mean_correct.append(correct.item() / (args.batch_size * args.npoint)) #计算每个batch的准确率
loss = criterion(seg_pred, target, trans_feat)
loss.backward()
optimizer.step()
train_instance_acc = np.mean(mean_correct) #计算每个epoch的准确率
log_string('Train accuracy is: %.5f' % train_instance_acc)
with torch.no_grad():
test_metrics = {}
total_correct = 0
total_seen = 0
total_seen_class = [0 for _ in range(num_part)]
total_correct_class = [0 for _ in range(num_part)]
shape_ious = {cat: [] for cat in seg_classes.keys()}
seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}
for cat in seg_classes.keys():
for label in seg_classes[cat]:
seg_label_to_cat[label] = cat
classifier = classifier.eval()
for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9):
cur_batch_size, NUM_POINT, _ = points.size()
points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda()
points = points.transpose(2, 1)
seg_pred, _ = classifier(points, to_categorical(label, num_classes)) #此时seg_pred的维度是[B,N,num_part]
cur_pred_val = seg_pred.cpu().data.numpy()
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) #此时cur_pred_val的维度是[B,N]
target = target.cpu().data.numpy() #此时target的维度是[B,N]
for i in range(cur_batch_size):
cat = seg_label_to_cat[target[i, 0]] #[i,0]只是为了获取批次中单样本的类别(cls)，那这就已经固定了点的seg_label的取值范围
logits = cur_pred_val_logits[i, :, :] #batch_size中的每个实例样本都是一个cls
cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0]#1-行方向搜索最大值，返回索引值;实现了提取预测的标签值
#for循环后cur_pred_val的维度是[B,N],其中的元素值为每个点对应的seg_label
correct = np.sum(cur_pred_val == target) #统计一个batch下预测与实际相同的个数
total_correct += correct #统计一个epoch下预测与实际相同的个数
total_seen += (cur_batch_size * NUM_POINT) #统计一个batch下所处理的点数
for l in range(num_part):
total_seen_class[l] += np.sum(target == l) #统计一个batch下是单个seg_label的点数，直到统计完一个epoch
total_correct_class[l] += (np.sum((cur_pred_val == l) & (target == l))) #np.sum只对True统计个数，统计完所有batch就完成了一个epoch
#for循环后，得到一个batch下每个seg_label预测与实际相同的个数，直到统计完一个epoch
for i in range(cur_batch_size):
segp = cur_pred_val[i, :] #获取预测的一个实例所有点的seg_label标签
segl = target[i, :] #获取实际的一个实例所有点的seg_label标签
cat = seg_label_to_cat[segl[0]] #获取该实例的类别，即获取一个cls
part_ious = [0.0 for _ in range(len(seg_classes[cat]))] #提醒此时的seg_classes[cat]是一个列表
for l in seg_classes[cat]:
if (np.sum(segl == l) == 0) and ( #实际和预测都不存在此seg_label，则置1
np.sum(segp == l) == 0): # part is not present, no prediction as well
part_ious[l - seg_classes[cat][0]] = 1.0
else: #&-交集，|-并集，获取part_ious
part_ious[l - seg_classes[cat][0]] = np.sum((segl == l) & (segp == l)) / float(
np.sum((segl == l) | (segp == l))) #cls有多少个部分，part_ious就有多少个元素
shape_ious[cat].append(np.mean(part_ious)) #针对每个cls，添加每个batch的MIou
#最终得到cls在每个batch下计算所得的MIou的列表集合
all_shape_ious = [] #一个epoch中有多少个batch，那么所属cls的列表就有多少个元素
for cat in shape_ious.keys():
for iou in shape_ious[cat]:
all_shape_ious.append(iou) #统计一个epoch中所有的seg_label的iou
shape_ious[cat] = np.mean(shape_ious[cat]) #计算一个cls的iou(该cls中seg_label的miou)
mean_shape_ious = np.mean(list(shape_ious.values())) #计算整体的iou(所有cls中iou的均值)
test_metrics['accuracy'] = total_correct / float(total_seen) #同时统计所有正确预测点占所有点的比例
test_metrics['class_avg_accuracy'] = np.mean(
np.array(total_correct_class) / np.array(total_seen_class, dtype=np.float)) #先计算每个seg_label的准确率，再取平均
for cat in sorted(shape_ious.keys()): #sorted()函数对所有可迭代的对象进行排序操作，因shape_ious是一个字典
log_string('eval mIoU of %s %f' % (cat + ' ' * (14 - len(cat)), shape_ious[cat]))
test_metrics['class_avg_iou'] = mean_shape_ious
test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious)
log_string('Epoch %d test Accuracy: %f Class avg mIOU: %f Inctance avg mIOU: %f' % (
epoch + 1, test_metrics['accuracy'], test_metrics['class_avg_iou'], test_metrics['inctance_avg_iou']))
if (test_metrics['inctance_avg_iou'] >= best_inctance_avg_iou):
logger.info('Save model...')
savepath = str(checkpoints_dir) + '/best_model.pth'
log_string('Saving at %s' % savepath)
state = {
'epoch': epoch,
'train_acc': train_instance_acc,
'test_acc': test_metrics['accuracy'],
'class_avg_iou': test_metrics['class_avg_iou'],
'inctance_avg_iou': test_metrics['inctance_avg_iou'],
'model_state_dict': classifier.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}
torch.save(state, savepath)
log_string('Saving model....')
if test_metrics['accuracy'] > best_acc:
best_acc = test_metrics['accuracy']
if test_metrics['class_avg_iou'] > best_class_avg_iou:
best_class_avg_iou = test_metrics['class_avg_iou']
if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou:
best_inctance_avg_iou = test_metrics['inctance_avg_iou']
log_string('Best accuracy is: %.5f' % best_acc)
log_string('Best class avg mIOU is: %.5f' % best_class_avg_iou)
log_string('Best inctance avg mIOU is: %.5f' % best_inctance_avg_iou)
global_epoch += 1
if __name__ == '__main__':
args = parse_args()
main(args)
【推荐】国内首个AI IDE，深度理解中文开发场景，立即下载体验Trae
【推荐】编程新体验，更懂你的AI，立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包，你的智能百科全书，全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell：AI 加持，快人一步
· winform 绘制太阳，地球，月球 运作规律
· AI与.NET技术实操系列（五）：向量存储与相似性搜索在 .NET 中的实现
· 超详细：普通电脑也行Windows部署deepseek R1训练数据并当服务器共享给他人
· 【硬核科普】Trae如何「偷看」你的代码？零基础破解AI编程运行原理
· 上周热点回顾（3.3-3.9）