def test_imagenet(self):
inputs = Variable(torch.randn(1,3,224,224))
f, params, stats = define_diracnet(18, 1, 'ImageNet')
outputs = f(inputs, params, stats, mode=False)
self.assertEqual(outputs.size(), torch.Size((1, 1000)))
def main():
opt = parser.parse_args()
print 'parsed options:', vars(opt)
epoch_step = json.loads(opt.epoch_step)
num_classes = 10 if opt.dataset == 'CIFAR10' else 100
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id
# to prevent opencv from initializing CUDA in workers
torch.randn(8).cuda()
os.environ['CUDA_VISIBLE_DEVICES'] = ''
train_loader = create_iterator(opt, True)
test_loader = create_iterator(opt, False)
f, params, stats = define_diracnet(opt.depth, opt.width, opt.dataset)
def create_optimizer(opt, lr):
print 'creating optimizer with lr = ', lr
return torch.optim.SGD(params.values(),
lr,
0.9,
weight_decay=opt.weightDecay)
optimizer = create_optimizer(opt, opt.lr)
epoch = 0
if opt.resume != '':
state_dict = torch.load(opt.resume)
epoch = state_dict['epoch']
params_tensors, stats = state_dict['params'], state_dict['stats']
for k, v in params.iteritems():
v.data.copy_(params_tensors[k])
optimizer.load_state_dict(state_dict['optimizer'])
print '\nParameters:'
kmax = max(len(key) for key in params.keys())
for i, (key, v) in enumerate(params.items()):
print str(i).ljust(5), key.ljust(kmax + 3), str(tuple(
v.size())).ljust(23), torch.typename(v.data)
print '\nAdditional buffers:'
kmax = max(len(key) for key in stats.keys())
for i, (key, v) in enumerate(stats.items()):
print str(i).ljust(5), key.ljust(kmax + 3), str(tuple(
v.size())).ljust(23), torch.typename(v)
n_parameters = sum(p.numel() for p in params.values())
print '\nTotal number of parameters:', n_parameters
meter_loss = tnt.meter.AverageValueMeter()
classacc = tnt.meter.ClassErrorMeter(topk=[1, 5], accuracy=True)
timer_train = tnt.meter.TimeMeter('s')
timer_test = tnt.meter.TimeMeter('s')
if not os.path.exists(opt.save):
os.mkdir(opt.save)
def h(sample):
inputs = Variable(cast(sample[0], opt.dtype))
targets = Variable(cast(sample[1], 'long'))
y = data_parallel(f, inputs, params, stats, sample[2],
np.arange(opt.ngpu))
return F.cross_entropy(y, targets), y
def log(t, state):
torch.save(
dict(params={k: v.data
for k, v in params.iteritems()},
stats=stats,
optimizer=state['optimizer'].state_dict(),
epoch=t['epoch']),
open(os.path.join(opt.save, 'model.pt7'), 'w'))
z = vars(opt).copy()
z.update(t)
print z
def on_sample(state):
state['sample'].append(state['train'])
def on_forward(state):
classacc.add(state['output'].data,
torch.LongTensor(state['sample'][1]))
meter_loss.add(state['loss'].data[0])
def on_start(state):
state['epoch'] = epoch
def on_start_epoch(state):
classacc.reset()
meter_loss.reset()
timer_train.reset()
state['iterator'] = tqdm(train_loader, ncols=100)
epoch = state['epoch'] + 1
if epoch in epoch_step:
lr = state['optimizer'].param_groups[0]['lr']
state['optimizer'] = create_optimizer(opt, lr * opt.lr_decay_ratio)
def on_end_epoch(state):
train_loss = meter_loss.value()
train_acc = classacc.value()
train_time = timer_train.value()
meter_loss.reset()
classacc.reset()
timer_test.reset()
engine.test(h, test_loader)
test_acc = classacc.value()
print log(
{
"train_loss": train_loss[0],
"train_acc": train_acc,
"test_loss": meter_loss.value()[0],
"test_acc": test_acc,
"epoch": state['epoch'],
"num_classes": num_classes,
"n_parameters": n_parameters,
"train_time": train_time,
"test_time": timer_test.value(),
}, state)
print '==> id: %s (%d/%d), test_acc: \33[91m%.2f\033[0m' % \
(opt.save, state['epoch'], opt.epochs, test_acc[0])
engine = Engine()
engine.hooks['on_sample'] = on_sample
engine.hooks['on_forward'] = on_forward
engine.hooks['on_start_epoch'] = on_start_epoch
engine.hooks['on_end_epoch'] = on_end_epoch
engine.hooks['on_start'] = on_start
engine.train(h, train_loader, opt.epochs, optimizer)
def test_cifar10(self):
inputs = Variable(torch.randn(1,3,32,32))
f, params, stats = define_diracnet(34, 1, 'CIFAR10')
outputs = f(inputs, params, stats, mode=False)
self.assertEqual(outputs.size(), torch.Size((1, 10)))
def main():
opt = parser.parse_args()
print('parsed options:', vars(opt))
epoch_step = json.loads(opt.epoch_step)
num_classes = 10 if opt.dataset == 'CIFAR10' else 100
torch.manual_seed(opt.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_id
def create_iterator(mode):
return DataLoader(create_dataset(opt, mode), opt.batch_size, shuffle=mode,
num_workers=opt.nthread, pin_memory=torch.cuda.is_available())
train_loader = create_iterator(True)
test_loader = create_iterator(False)
f, params = define_diracnet(opt.depth, opt.width, opt.dataset)
def create_optimizer(opt, lr):
print('creating optimizer with lr = ', lr)
params_wd, params_rest = [], []
for k, v in params.items():
if v.requires_grad:
(params_wd if v.dim() != 1 else params_rest).append(v)
groups = [{'params': params_wd, 'weight_decay': opt.weight_decay},
{'params': params_rest}]
return SGD(groups, lr, momentum=0.9)
optimizer = create_optimizer(opt, opt.lr)
epoch = 0
if opt.resume != '':
state_dict = torch.load(opt.resume)
epoch = state_dict['epoch']
params_tensors = state_dict['params']
for k, v in params.items():
v.data.copy_(params_tensors[k])
optimizer.load_state_dict(state_dict['optimizer'])
print('\nParameters:')
print_tensor_dict(params)
n_parameters = sum(p.numel() for p in params.values() if p.requires_grad)
print('\nTotal number of parameters:', n_parameters)
meter_loss = tnt.meter.AverageValueMeter()
classacc = tnt.meter.ClassErrorMeter(accuracy=True)
timer_train = tnt.meter.TimeMeter('s')
timer_test = tnt.meter.TimeMeter('s')
if not os.path.exists(opt.save):
os.mkdir(opt.save)
def h(sample):
inputs = cast(sample[0], opt.dtype)
targets = cast(sample[1], 'long')
y = data_parallel(f, inputs, params, sample[2], range(opt.ngpu)).float()
return F.cross_entropy(y, targets), y
def log(t, state):
torch.save(dict(params=params, epoch=t['epoch'], optimizer=state['optimizer'].state_dict()),
os.path.join(opt.save, 'model.pt7'))
z = {**vars(opt), **t}
with open(os.path.join(opt.save, 'log.txt'), 'a') as flog:
flog.write('json_stats: ' + json.dumps(z) + '\n')
print(z)
def on_sample(state):
state['sample'].append(state['train'])
def on_forward(state):
loss = float(state['loss'])
classacc.add(state['output'].data, state['sample'][1])
meter_loss.add(loss)
if state['train']:
state['iterator'].set_postfix(loss=loss)
def on_start(state):
state['epoch'] = epoch
def on_start_epoch(state):
classacc.reset()
meter_loss.reset()
timer_train.reset()
state['iterator'] = tqdm(train_loader, dynamic_ncols=True)
epoch = state['epoch'] + 1
if epoch in epoch_step:
lr = state['optimizer'].param_groups[0]['lr']
state['optimizer'] = create_optimizer(opt, lr * opt.lr_decay_ratio)
def on_end_epoch(state):
train_loss = meter_loss.value()
train_acc = classacc.value()
train_time = timer_train.value()
meter_loss.reset()
classacc.reset()
timer_test.reset()
with torch.no_grad():
engine.test(h, test_loader)
test_acc = classacc.value()[0]
print(log({
"train_loss": train_loss[0],
"train_acc": train_acc[0],
"test_loss": meter_loss.value()[0],
"test_acc": test_acc,
"epoch": state['epoch'],
"num_classes": num_classes,
"n_parameters": n_parameters,
"train_time": train_time,
"test_time": timer_test.value(),
}, state))
print('==> id: %s (%d/%d), test_acc: \33[91m%.2f\033[0m' %
(opt.save, state['epoch'], opt.epochs, test_acc))
engine = Engine()
engine.hooks['on_sample'] = on_sample
engine.hooks['on_forward'] = on_forward
engine.hooks['on_start_epoch'] = on_start_epoch
engine.hooks['on_end_epoch'] = on_end_epoch
engine.hooks['on_start'] = on_start
engine.train(h, train_loader, opt.epochs, optimizer)