model = ResUNet.Res16UNet34C(3, num_class).to(device) evaluator = Evaluator(num_class) model_dict = torch.load(os.path.join(config['resume_path'], 'parameters.pth'))["model"] model.load_state_dict(model_dict) model.eval() test_data = read_plyfile("/home/gaoqiyu/文档/Stanford3dDataset_v1.2_Aligned_Version/ply/val/Area_1_hallway_2.ply") sinput, coords, feats, labels = data_preprocess(test_data, voxel_size) soutput = model(sinput.to(device)) # Feed-forward pass and get the prediction _, pred = soutput.F.max(1) pred = pred.cpu().numpy() evaluator.generate(pred, labels) IOU, mIOU = evaluator.mIOU() print("IOU: ", IOU) print("mIOU: ", mIOU) # Map color pred_pcd = o3d.geometry.PointCloud() coordinates = soutput.C.numpy()[:, :3] # last c colors_pred = np.array([SCANNET_COLOR_MAP[VALID_CLASS_IDS[l]] for l in pred]) # Create a point cloud fileolumn is the batch index pred_pcd.points = o3d.utility.Vector3dVector(coordinates * voxel_size) pred_pcd.colors = o3d.utility.Vector3dVector(colors_pred / 255) colors_ground = np.array([SCANNET_COLOR_MAP[VALID_CLASS_IDS[l]] for l in labels]) ground_pcd = o3d.geometry.PointCloud() ground_pcd.points = o3d.utility.Vector3dVector(coordinates * voxel_size + np.array([0, 10, 0])) ground_pcd.colors = o3d.utility.Vector3dVector(colors_ground / 255)
class Trainer(object): def __init__(self, config_): self.config = config_ self.best_pred = -math.inf self.train_iter_number = 0 self.val_iter_number = 0 self.epoch = 0 self.class_name = self.config["class_label"] if self.config["multi_gpu"]: self.device_list = list(range(torch.cuda.device_count())) self.device = self.device_list[0] else: self.device = torch.device('cuda') self.loss_value = torch.tensor(0.0, requires_grad=True).to(self.device) self.point_number = self.config["point_num"] self.batch_size = self.config["batch_size"] self.model = ResUNet.Res16UNet34C(3, self.config["class"]) if self.config["fine_tune"]: model_dict = torch.load( os.path.join(self.config["resume_path"], 'weights_14.pth'), map_location=lambda storage, loc: storage.cuda(self.device)) self.model.load_state_dict(model_dict) self.optimizer = torch.optim.SGD( [{ 'params': self.model.convtr7p2s2.parameters(), 'lr': self.config["lr"] / 1e2 }, { 'params': self.model.bntr7.parameters(), 'lr': self.config["lr"] / 1e2 }, { 'params': self.model.block8.parameters(), 'lr': self.config["lr"] / 1e1 }, { 'params': self.model.final.parameters(), 'lr': self.config["lr"] }], lr=self.config["lr"] / 1e4, momentum=self.config["momentum"], weight_decay=1e-4) if self.config["use_cuda"]: self.model = self.model.to(self.device) if self.config["multi_gpu"]: self.model = torch.nn.DataParallel(self.model, device_ids=self.device_list) self.loss = torch.nn.CrossEntropyLoss( ignore_index=self.config['ignore_label']) self.train_data = initialize_data_loader(S3DISDataset, self.config, phase='TRAIN', threads=1, augment_data=True, shuffle=True, repeat=True, batch_size=1, limit_numpoints=False) self.val_data = initialize_data_loader(S3DISDataset, self.config, threads=1, phase='VAL', augment_data=False, shuffle=True, repeat=False, batch_size=1, limit_numpoints=False) self.optimizer = torch.optim.SGD( self.model.parameters(), lr=self.config['lr'], momentum=self.config['momentum'], weight_decay=self.config['weight_decay']) # self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config['lr'], weight_decay=self.config['weight_decay']) self.lr_scheduler = PolyLR(self.optimizer, max_iter=60000, power=self.config['poly_power'], last_step=-1) log_path = os.path.join(self.config["log_path"], str(time.time())) os.mkdir(log_path) if not os.path.exists(log_path) else None self.summary = SummaryWriter(log_path) self.evaluator = Evaluator(self.config["class"]) self.load() def train(self, epoch_): epoch_loss = [] self.model.train() for ith, data_dict in enumerate(self.train_data): point, labels = self.data_preprocess(data_dict, 'train') output_sparse = self.model(point) pred = output_sparse.F self.reset_loss(labels) self.loss_value = self.loss(pred, labels) + self.loss_value epoch_loss.append( self.loss(pred, labels).item() / self.config["accumulate_gradient"]) self.train_iter_number += 1 if self.train_iter_number % self.batch_size == 0: self.loss_value /= (self.config["accumulate_gradient"] * self.batch_size) self.loss_value.backward() lr_value = self.optimizer.state_dict()['param_groups'][0]['lr'] self.optimizer.step() self.lr_scheduler.step(self.train_iter_number) self.optimizer.zero_grad() torch.cuda.empty_cache() self.summary.add_scalar('train/loss: ', self.loss_value, self.train_iter_number) self.summary.add_scalar('train/lr: ', lr_value, self.train_iter_number) print( "train epoch: {}/{}, ith: {}/{}, loss: {:.4f}, lr: {:.6f}" .format(epoch_, self.config['epoch'], self.train_iter_number, len(self.train_data), self.loss_value.item(), lr_value)) self.loss_value = 0 if ith == len(self.train_data) - 1: break average_loss = np.nanmean(epoch_loss) self.summary.add_scalar('train/loss_epoch: ', average_loss, epoch_) print("epoch: {}/{}, average_loss: {:.4f}".format( epoch_, self.config['epoch'], average_loss)) print( '------------------------------------------------------------------' ) def eval(self, epoch_): self.model.eval() torch.cuda.empty_cache() IOU_epoch = [] mIOU_epoch = [] Acc_epoch = [] precision_epoch = [] recall_epoch = [] epoch_loss = [] for ith, data_dict in enumerate(self.val_data): point, labels = self.data_preprocess(data_dict, 'val') with torch.no_grad(): output = self.model(point) pred = output.F loss_eval = self.loss(pred, labels) / self.config["accumulate_gradient"] epoch_loss.append(loss_eval.item()) mIOU, Acc, precision, recall = self.evaluator.generate( pred.max(1)[1].cpu(), labels.cpu()) IOU, _ = self.evaluator.mIOU() IOU_epoch.append(IOU) mIOU_epoch.append(mIOU) Acc_epoch.append(Acc) precision_epoch.append(precision) recall_epoch.append(recall) self.val_iter_number += 1 self.summary.add_scalar('val/mIOU', mIOU, self.val_iter_number) self.summary.add_scalar('val/accuracy', Acc, self.val_iter_number) self.summary.add_scalar('val/precision', precision, self.val_iter_number) self.summary.add_scalar('val/recall', recall, self.val_iter_number) self.summary.add_scalar('val/loss: ', loss_eval, self.val_iter_number) print( "val epoch: {}/{}, ith: {}/{}, loss: {:.4f}, mIOU: {:.2%}, accuracy {:.2%},precision {:.2%},recall {:.2%}" .format(epoch_, self.config['epoch'], ith, len(self.val_data), loss_eval, mIOU, Acc, precision, recall)) average_loss = np.nanmean(epoch_loss) average_IOU = np.nanmean(IOU_epoch, axis=0) average_mIOU = np.nanmean(mIOU_epoch) average_Acc = np.nanmean(Acc_epoch) average_precision = np.nanmean(precision_epoch) average_recall = np.nanmean(recall_epoch) self.summary.add_scalar('val/acc_epoch', average_Acc, epoch_) self.summary.add_scalar('val/loss_epoch', average_loss, epoch_) self.summary.add_scalar('val/mIOU_epoch', average_mIOU, epoch_) self.summary.add_scalar('val/precision_epoch', average_precision, epoch_) self.summary.add_scalar('val/recall_epoch', average_recall, epoch_) print( "epoch: {}/{}, average_loss: {:.4f},average_mIOU: {:.2%}, average_accuracy: {:.2%}, average_precision: {:.2%}, average_recall: {:.2%}" .format(epoch_, self.config['epoch'], average_loss, average_mIOU, average_Acc, average_precision, average_recall)) print("class name ", self.class_name) print("IOU/class: ", average_IOU) print( '------------------------------------------------------------------' ) if average_mIOU > self.best_pred: self.best_pred = average_mIOU self.save(epoch_) def load(self): load_path = os.path.join(self.config["resume_path"], 'parameters.pth') if os.path.isfile(load_path): load_parameters = torch.load( load_path, map_location=lambda storage, loc: storage.cuda(self.device)) self.optimizer.load_state_dict(load_parameters['optimizer']) self.lr_scheduler.load_state_dict(load_parameters['lr_scheduler']) self.model.load_state_dict(load_parameters['model']) self.best_pred = load_parameters['best_prediction'] self.epoch = load_parameters['epoch'] self.train_iter_number = load_parameters['train_iter_number'] self.val_iter_number = load_parameters['val_iter_number'] self.loss_value = load_parameters['loss_value'] self.model = self.model.to(self.device) self.lr_scheduler.step(self.train_iter_number) def save(self, epoch_): os.mkdir(self.config["resume_path"]) if not os.path.exists( self.config["resume_path"]) else None torch.save( { 'best_prediction': self.best_pred, 'model': self.model.state_dict(), 'optimizer': self.optimizer.state_dict(), 'lr_scheduler': self.lr_scheduler.state_dict(), 'epoch': epoch_, 'train_iter_number': self.train_iter_number, 'val_iter_number': self.val_iter_number, 'loss_value': self.loss_value }, os.path.join(self.config["resume_path"], 'parameters.pth')) def data_preprocess(self, data_dict, data_type='train'): coords = data_dict[0] feats = data_dict[1] / 255 - 0.5 labels = data_dict[2] length = coords.shape[0] if length > self.point_number and data_type == 'train': inds = np.random.choice(np.arange(length), self.point_number, replace=False) coords, feats, labels = coords[inds], feats[inds], labels[inds] if data_type == 'train': # For some networks, making the network invariant to even, odd coords is important coords[:, :3] += (torch.rand(3) * 100).type_as(coords) points = ME.SparseTensor(feats, coords.int()) if self.config["use_cuda"]: points, labels = points.to(self.device), labels.to(self.device) return points, labels.long() def reset_loss(self, label): data_ = label.cpu().numpy() number = data_.shape[0] bin = np.bincount(data_, minlength=self.config["class"]) for ith, i in enumerate(bin): if i != 0: bin[ith] = number / i bin = bin / np.sum(bin) self.loss = torch.nn.CrossEntropyLoss( weight=torch.from_numpy(bin).cuda().float(), ignore_index=self.config['ignore_label'])