Commit 2d231c2b authored by Kai Chen's avatar Kai Chen
Browse files

bug fix

parent 7478e1a6
...@@ -5,17 +5,14 @@ from .hook import Hook ...@@ -5,17 +5,14 @@ from .hook import Hook
class OptimizerHook(Hook): class OptimizerHook(Hook):
def __init__(self, grad_clip=False, max_norm=35, norm_type=2): def __init__(self, grad_clip=None):
self.grad_clip = grad_clip self.grad_clip = grad_clip
self.max_norm = max_norm
self.norm_type = norm_type
def after_train_iter(self, runner): def after_train_iter(self, runner):
runner.optimizer.zero_grad() runner.optimizer.zero_grad()
runner.outputs['loss'].backward() runner.outputs['loss'].backward()
if self.grad_clip: if self.grad_clip is not None:
clip_grad.clip_grad_norm_( clip_grad.clip_grad_norm_(
filter(lambda p: p.requires_grad, runner.model.parameters()), filter(lambda p: p.requires_grad, runner.model.parameters()),
max_norm=self.max_norm, **self.grad_clip)
norm_type=self.norm_type)
runner.optimizer.step() runner.optimizer.step()
...@@ -2,6 +2,7 @@ import multiprocessing ...@@ -2,6 +2,7 @@ import multiprocessing
import torch import torch
import mmcv
from .io import load_checkpoint from .io import load_checkpoint
...@@ -61,8 +62,7 @@ def parallel_test(model_cls, ...@@ -61,8 +62,7 @@ def parallel_test(model_cls,
idx_queue.put(i) idx_queue.put(i)
results = [None for _ in range(len(dataset))] results = [None for _ in range(len(dataset))]
import cvbase as cvb prog_bar = mmcv.ProgressBar(task_num=len(dataset))
prog_bar = cvb.ProgressBar(task_num=len(dataset))
for _ in range(len(dataset)): for _ in range(len(dataset)):
idx, res = result_queue.get() idx, res = result_queue.get()
results[idx] = res results[idx] = res
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment