Skip to content

Commit ef62758

Browse files
changed default train config
1 parent 4055e30 commit ef62758

File tree

2 files changed

+5
-3
lines changed

2 files changed

+5
-3
lines changed

examples/pytorch/FastCells/train_classifier.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ def configure_optimizer(self, options):
106106
step_sizes=oo.step_sizes)
107107
elif options.optimizer == "SGD":
108108
optimizer = optim.SGD(self.parameters(), lr=initial_rate, weight_decay=oo.weight_decay,
109-
momentum=oo.momentum, dampening=oo.dampening)
109+
momentum=oo.momentum, dampening=oo.dampening, nesterov=oo.nesterov)
110110
return optimizer
111111

112112
def configure_lr(self, options, optimizer, ticks, total_iterations):
@@ -478,7 +478,8 @@ def train(config, evaluate_only=False, outdir=".", detail=False, azureml=False):
478478
model.cuda() # move the processing to GPU
479479

480480
start = time.time()
481-
log = model.fit(training_data, validation_data, config.training, config.model.sparsify, device, detail, run)
481+
log = model.fit(training_data, validation_data, config.training,
482+
config.model.sparsify, device, detail, run)
482483
end = time.time()
483484

484485
passed, total, rate = model.evaluate(training_data, batch_size, device)

examples/pytorch/FastCells/training_config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ def __init__(self):
4848
self.etas = (0.5, 1.2) # Rprop
4949
self.dampening = 0 # SGD
5050
self.step_sizes = (1e-06, 50) # Rprop
51+
self.nesterov = True # SGD
5152

5253

5354
class TrainingOptions:
@@ -60,7 +61,7 @@ def __init__(self):
6061
self.lr_gamma = 1
6162
self.lr_step_size = 1
6263
self.batch_size = 128
63-
self.optimizer = "Adam"
64+
self.optimizer = "SGD"
6465
self.optimizer_options = OptimizerOptions()
6566
self.use_gpu = False
6667
self.rolling = False

0 commit comments

Comments
 (0)