def __init__(self,
cfg,
max_iter=None,
batch_size=None,
train_devices=None,
model_save_step=None,
model_save_root=None,
vis=None,
vis_step=None):
"""
训练器初始化
值为None的参数项不指定时为默认,已在配置文件中设置.
如需更改参数建议在Configs配置文件中进行更改
不建议直接指定参数,只留做扩展用.
:param cfg: 配置
:param max_iter: 最大训练轮数
:param batch_size: 批次数,
:param train_devices: 训练设备,列表,eg:[0,1],使用0,1俩个GPU,这里0,1为gpu编号,可用nvidia-smi查看.,不指定时为默认,已在配置文件中设置
:param vis: visdom.Visdom(),用于训练过程可视化.绘制损失曲线已经学习率
:param model_save_step: 模型保存步长
:param vis_step: visdom可视化步长
"""
self.cfg = cfg
self.iterations = self.cfg.TRAIN.MAX_ITER
if max_iter:
self.iterations = max_iter
self.batch_size = cfg.TRAIN.BATCH_SIZE
if batch_size:
self.batch_size = batch_size
self.train_devices = cfg.DEVICE.TRAIN_DEVICES
if train_devices:
self.train_devices = train_devices
self.model_save_root = cfg.FILE.MODEL_SAVE_ROOT
if model_save_root:
self.model_save_root = model_save_root
if not os.path.exists(self.model_save_root):
os.mkdir(self.model_save_root)
self.model_save_step = self.cfg.STEP.MODEL_SAVE_STEP
if model_save_step:
self.model_save_step = model_save_step
self.vis = setup_visdom()
if vis:
self.vis = vis
self.vis_step = self.cfg.STEP.VIS_STEP
if vis_step:
self.vis_step = vis_step
self.model = None
self.loss_func = None
self.optimizer = None
self.scheduler = None
# -*- coding: utf-8 -*-
# @Author : LG
from Data.Dataset import indoor3d_Dataset
from torch.utils.data import DataLoader
from Model.Pointnet2 import PointnetMSG
from torch import nn
import torch
from torch.nn import DataParallel
from Utils.visdom_op import setup_visdom, visdom_line
EPOCH = 100
LR = 0.0001
vis = setup_visdom()
train_data = indoor3d_Dataset(is_train=True, data_root='Data', test_area=5)
train_loader = DataLoader(train_data,
batch_size=32,
shuffle=True,
num_workers=4)
model = PointnetMSG(xyz_channel=3, data_channel=6, num_classes=13)
model.load_state_dict(
torch.load(
'/home/super/PycharmProjects/pointnet2/Weights/model_40000.pkl'))
model = model.to('cuda')
model = DataParallel(model, device_ids=[0, 1])
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)