def Train(cfg):
os.environ['CUDA_DEVICES_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
[str(item) for item in cfg.TASK.GPUS])
device = torch.device('cuda' if len(cfg.TASK.GPUS) > 0 else 'cpu')
# init logger
output_dir = os.path.join(cfg.TASK.OUTPUT_ROOT_DIR, cfg.TASK.NAME)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger(cfg.TASK.NAME, output_dir, distributed_rank=0)
# data loader
train_loader = MakeTrainLoader(
os.path.join(cfg.DATA.ROOT_DIR, cfg.DATA.TRAIN_PATH),
cfg.DATA.TRAIN_BATCH)
num_images = len(train_loader.dataset)
print('total train data: ', num_images)
# model
model = Net().to(device)
# loss
loss_funcs = nn.CrossEntropyLoss().to(device)
# optimizer
params = [p for n, p in model.named_parameters()]
param_groups = [{'params': params, 'lr': 0.1}]
# optimizer = optim.SGD(param_groups, momentum=0.9, weight_decay=5e-4)
optimizer = optim.Adam(param_groups)
# lr scheduler
lr_scheduler = BuildLRScheduler(optimizer, [50, 100, 200], 0.1)
# start train
for epoch_idx in range(0, cfg.SOLVER.EPOCHS):
logger.info('train epoch: {0}'.format(epoch_idx))
TrainEpoch(epoch_idx, train_loader, model, loss_funcs, optimizer,
lr_scheduler, device, logger)
N = int(config["num_epoch"])
auto_save = int(config["auto_save"])
start_epoch = 0
opt_name = config["opt"]
### model and reader
net_model = Net(config)
net_reader = NetReader(config)
net_model.train(mode=True)
criterion = net_model.get_criterion(config)
data_loader = net_reader.getDataLoader()
### trainable and restorable
trainable_var = OrderedDict(net_model.get_trainable())
untrainable_var = OrderedDict(net_model.named_parameters())
for key, val in trainable_var.items():
del untrainable_var[key]
restore_var = OrderedDict(net_model.get_restorable())
loss = 0
### optimizer detection
opt_params = [
{
'params': list(trainable_var.values()),
'lr': lr,
},
{