def main(args):
#------------ start to prepare dataset ------------'
tr_dataset = Dataset(list_dir=args.train_dir, cv=0)
cv_dataset = Dataset(list_dir=args.valid_dir, cv=1)
tr_loader = DataLoader(tr_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
cv_loader = DataLoader(cv_dataset,
batch_size=2,
shuffle=False,
num_workers=0)
#'------------------ model -----------------------'
model = Model(kernel_size=3, stride=1, dropout=0.1)
print(model)
model.apply(weight_init)
if args.use_cuda == True and torch.cuda.is_available():
device = torch.device("cuda")
model = torch.nn.DataParallel(model)
else:
device = torch.device('cpu')
model = model.to(device=device)
# optimizer
if args.optimizer == 'RMSprop':
optimizier = torch.optim.RMSprop(model.parameters(), lr=args.lr)
elif args.optimizer == 'Adam':
optimizier = torch.optim.Adam(model.parameters(), lr=args.lr)
else:
print("Not support optimizer")
return RuntimeError('Unrecognized optimizer')
# Loss
# Loss = torch.nn.MSELoss()
train_total_loss = []
cv_total_loss = []
best_loss = float("inf")
no_improve_nums = 0
# ---------------------------------- Training ------------------------
for epoch in range(0, args.epochs):
model.train()
tr_loss = torch.tensor(0.0)
for i, (data) in enumerate(tr_loader):
x, y = data
x = x.to(device=device, dtype=torch.float32)
y = y.to(device=device, dtype=torch.long)
est = model(x)
loss = torch.nn.functional.cross_entropy(input=est, target=y)
# loss = Loss(input=est, target=y)
tr_loss += loss
optimizier.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(parameters=model.parameters(),
max_norm=5)
optimizier.step()
tr_loss = tr_loss / i
train_total_loss.append(tr_loss.cpu().detach().numpy())
print('-' * 80)
print('Epoch %d End train with loss: %.3f' % (epoch, tr_loss))
print('-' * 80)
# ---------------------------- validation ---------------------------
model.eval()
cv_loss = torch.tensor(0.0)
with torch.no_grad():
for j, (data) in enumerate(cv_loader):
x, y = data
x = x.to(device=device, dtype=torch.float)
y = y.to(device=device, dtype=torch.long)
est = model(x)
loss = torch.nn.functional.cross_entropy(input=est, target=y)
# loss = Loss(input=est, target=y)
cv_loss += loss
if j % 5 == 0:
print('Epoch %d, Iter: %d, Loss: %.3f' % (epoch, j, loss))
cv_loss = cv_loss / j
cv_total_loss.append(cv_loss.cpu().detach().numpy())
print('-' * 80)
if best_loss > cv_loss:
best_loss = cv_loss
torch.save(
model.module.serialize(model.module,
optimizier,
epoch + 1,
tr_loss=tr_loss,
cv_loss=cv_loss),
args.save_folder / args.save_name)
print("Find best validation model, saving to %s" %
str(args.save_folder / args.save_name))
no_improve_nums = 0
else:
print('no improve ...')
no_improve_nums += 1
if no_improve_nums >= 3:
optim_state = optimizier.state_dict()
optim_state['param_groups'][0][
'lr'] = optim_state['param_groups'][0]['lr'] / 2.0
optimizier.load_state_dict(optim_state)
print('Reduce learning rate to lr: %.8f' %
optim_state['param_groups'][0]['lr'])
if no_improve_nums >= 6:
print('No improve for 6 epochs, stopping')
break
print('Epoch %d End validation with loss: %.3f, best loss: %.3f' %
(epoch, cv_loss, best_loss))
print('-' * 80)
writer = SummaryWriter()
device = torch.device('cuda:0')
model = Model()
model = model.to(device)
imsize = 256, 256
print(f"TrainSize = {len(train)}")
print(f"Raw1Size = {len(raw1)}")
print(f"Raw2Size = {len(raw2)}")
print(f"Raw3Size = {len(raw3)}")
optimizer = torch.optim.Adam(model.parameters(), lr=1.0e-5)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 4, 0.8)
model.train()
for epoch in range(50):
print(f"------ EPOCH {epoch} ------")
# -------------------------- TRAIN --------------------------
train_losses = []
for i_batch, sample_batched in tqdm(enumerate(train_loader)):
optimizer.zero_grad()
pred_images, next_images, event_images, _ = sample_batched
pred_images = pred_images.to(device)
next_images = next_images.to(device)
event_images = event_images.to(device)
flow = model.forward(event_images)
