def train(args):
# parse config
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
config = parse_config(args.config)
train_config = merge_configs(config, 'train', vars(args))
#根据自己定义的网络，声明train_model
train_model = TSN(args.batch_size, 32)
train_model.train()
opt = fluid.optimizer.Momentum(0.001,
0.9,
parameter_list=train_model.parameters())
if args.pretrain:
# 加载上一次训练的模型，继续训练
model, _ = fluid.dygraph.load_dygraph(args.save_dir + '/tsn_model')
train_model.load_dict(model)
# build model
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
# 迭代器，每次输送一个批次的数据，每个数据的形式(imgs, label)
train_reader = UCFReader('train', train_config).create_reader()
epochs = args.epoch
for i in range(epochs):
for batch_id, data in enumerate(train_reader()):
dy_x_data = np.array([x[0] for x in data]).astype(
'float32'
) # dy_x_data=imgs shape:(batch_size, seg_num*seglen, 3, target_size, target_size)
y_data = np.array([x[1] for x in data]).astype(
'int64') # y_data=label shape:(batch_size,)
img = fluid.dygraph.to_variable(dy_x_data)
label = fluid.dygraph.to_variable(y_data)
label = fluid.layers.reshape(label, [args.batch_size, 1])
label.stop_gradient = True # 反向求导时，只有目标函数中的label部分对label求导（而不是label与img融合求导）
out, acc = train_model(img, label)
loss = fluid.layers.cross_entropy(out, label)
avg_loss = fluid.layers.mean(loss)
avg_loss.backward()
opt.minimize(avg_loss)
train_model.clear_gradients()
if batch_id % 100 == 0:
print("Loss at epoch {} step {}: {}, acc: {}".format(
i, batch_id, avg_loss.numpy(), acc.numpy()))
fluid.dygraph.save_dygraph(train_model.state_dict(),
args.save_dir + '/tsn_model')
print("Final loss: {}".format(avg_loss.numpy()))
def main():
global args, best_prec1
args = parser.parse_args()
check_rootfolders()
categories, args.train_list, args.val_list, args.root_path, prefix = datasets_video.return_dataset(
args.dataset, args.modality)
num_class = len(categories)
args.store_name = '_'.join([
'TRN', args.dataset, args.modality, args.arch, args.consensus_type,
'segment%d' % args.num_segments
])
print('storing name: ' + args.store_name)
model = TSN(339,
args.num_segments,
args.modality,
base_model=args.arch,
consensus_type=args.consensus_type,
dropout=args.dropout,
img_feature_dim=args.img_feature_dim,
partial_bn=not args.no_partialbn)
# freeze cnn
# _, cnn =list(model.named_children())[0]
# for p in cnn.parameters():
# p.requires_grad = False
crop_size = model.crop_size
scale_size = model.scale_size
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
train_augmentation = model.get_augmentation()
model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()
# remove if not transfer learning
checkpoint = torch.load('/home/ec2-user/mit_weights.pth.tar')
model.load_state_dict(checkpoint['state_dict'])
# only if transfer learning
# for module in list(list(model._modules['module'].children())[-1].children())[-1].children():
# module[-1] = nn.Linear(256, num_class)
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
model.cuda()
model.train(True)
cudnn.benchmark = True
# Data loading code
if args.modality != 'RGBDiff':
normalize = GroupNormalize(input_mean, input_std)
else:
normalize = IdentityTransform()
if args.modality == 'RGB':
data_length = 1
elif args.modality in ['Flow', 'RGBDiff']:
data_length = 5
train_loader = torch.utils.data.DataLoader(TSNDataSet(
args.root_path,
args.train_list,
num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl=prefix,
transform=torchvision.transforms.Compose([
train_augmentation,
Stack(roll=(args.arch in ['BNInception', 'InceptionV3'])),
ToTorchFormatTensor(
div=(args.arch not in ['BNInception', 'InceptionV3'])),
normalize,
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(TSNDataSet(
args.root_path,
args.val_list,
num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl=prefix,
random_shift=False,
transform=torchvision.transforms.Compose([
GroupScale(int(scale_size)),
GroupCenterCrop(crop_size),
Stack(roll=(args.arch in ['BNInception', 'InceptionV3'])),
ToTorchFormatTensor(
div=(args.arch not in ['BNInception', 'InceptionV3'])),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# define loss function (criterion) and optimizer
if args.loss_type == 'nll':
criterion = torch.nn.CrossEntropyLoss().cuda()
else:
raise ValueError("Unknown loss type")
for group in policies:
print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
group['name'], len(group['params']), group['lr_mult'],
group['decay_mult'])))
optimizer = torch.optim.Adam(model.parameters(), args.lr)
# optimizer = torch.optim.SGD(policies,
# args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
if args.evaluate:
validate(val_loader, model, criterion, 0)
return
log_training = open(
os.path.join(args.root_log, '%s.csv' % args.store_name), 'w')
for epoch in range(args.start_epoch, args.epochs):
# adjust_learning_rate(optimizer, epoch, args.lr_steps)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, log_training)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
prec1 = validate(val_loader, model, criterion,
(epoch + 1) * len(train_loader), log_training,
epoch)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
summary_writer.close()
