allow arbitrary batch size
@@ -168,9 +168,9 @@ def solve_system(ode_system, conditions, t_min, t_max,
criterion = nn.MSELoss()
n_examples = example_generator.size
- if n_examples%batch_size != 0:
- raise RuntimeError('Please choose a batch_size such that it is a factor of the size of the training set.')
- n_batches = n_examples//batch_size
+# if n_examples%batch_size != 0:
+# raise RuntimeError('Please choose a batch_size such that it is a factor of the size of the training set.')
+# n_batches = n_examples//batch_size
zeros = torch.zeros(batch_size)
loss_history = []
@@ -179,8 +179,13 @@ def solve_system(ode_system, conditions, t_min, t_max,
loss_epoch = 0.0
examples = example_generator.get_examples()
- ts_batches = examples.reshape((n_batches, batch_size, 1))
- for ts in ts_batches:
+# ts_batches = examples.reshape((n_batches, batch_size, 1))
+ examples = examples.reshape(n_examples, 1)
+ batch_start, batch_end = 0, batch_size
+ while batch_start < n_examples:
+
+ if batch_end >= n_examples: batch_end = n_examples
+ ts = examples[batch_start:batch_end]
# the dependent variables
vs = []
@@ -192,13 +197,16 @@ def solve_system(ode_system, conditions, t_min, t_max,
Fvts = ode_system(*vs, ts)
loss = 0.0
for Fvt in Fvts: loss += criterion(Fvt, zeros)
- loss_epoch += loss.item()
+ loss_epoch += loss.item()/(batch_end-batch_start) # assume the loss is a mean over all examples
optimizer.zero_grad()
loss.backward()
optimizer.step()
- loss_history.append(loss_epoch/n_batches)
+ batch_start += batch_size
+ batch_end += batch_size
+
+ loss_history.append(loss_epoch)
if loss_history[-1] < tol: break
if monitor and epoch%monitor.check_every == 0:
