| 165 | # Ray Tune integration points: |
| 166 | |
| 167 | def train_cifar(config, data_dir=None): |
| 168 | net = Net(config["l1"], config["l2"]) |
| 169 | device = config["device"] |
| 170 | |
| 171 | net = net.to(device) |
| 172 | if torch.cuda.device_count() > 1: |
| 173 | net = nn.DataParallel(net) |
| 174 | |
| 175 | criterion = nn.CrossEntropyLoss() |
| 176 | optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9) |
| 177 | |
| 178 | # Load checkpoint if resuming training |
| 179 | checkpoint = tune.get_checkpoint() |
| 180 | if checkpoint: |
| 181 | with checkpoint.as_directory() as checkpoint_dir: |
| 182 | checkpoint_path = Path(checkpoint_dir) / "checkpoint.pt" |
| 183 | checkpoint_state = torch.load(checkpoint_path) |
| 184 | start_epoch = checkpoint_state["epoch"] |
| 185 | net.load_state_dict(checkpoint_state["net_state_dict"]) |
| 186 | optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"]) |
| 187 | else: |
| 188 | start_epoch = 0 |
| 189 | |
| 190 | trainset, _testset = load_data(data_dir) |
| 191 | |
| 192 | test_abs = int(len(trainset) * 0.8) |
| 193 | train_subset, val_subset = random_split( |
| 194 | trainset, [test_abs, len(trainset) - test_abs] |
| 195 | ) |
| 196 | |
| 197 | trainloader = torch.utils.data.DataLoader( |
| 198 | train_subset, batch_size=int(config["batch_size"]), shuffle=True, num_workers=8 |
| 199 | ) |
| 200 | valloader = torch.utils.data.DataLoader( |
| 201 | val_subset, batch_size=int(config["batch_size"]), shuffle=True, num_workers=8 |
| 202 | ) |
| 203 | |
| 204 | for epoch in range(start_epoch, 10): # loop over the dataset multiple times |
| 205 | running_loss = 0.0 |
| 206 | epoch_steps = 0 |
| 207 | for i, data in enumerate(trainloader, 0): |
| 208 | # get the inputs; data is a list of [inputs, labels] |
| 209 | inputs, labels = data |
| 210 | inputs, labels = inputs.to(device), labels.to(device) |
| 211 | |
| 212 | # zero the parameter gradients |
| 213 | optimizer.zero_grad() |
| 214 | |
| 215 | # forward + backward + optimize |
| 216 | outputs = net(inputs) |
| 217 | loss = criterion(outputs, labels) |
| 218 | loss.backward() |
| 219 | optimizer.step() |
| 220 | |
| 221 | # print statistics |
| 222 | running_loss += loss.item() |
| 223 | epoch_steps += 1 |
| 224 | if i % 2000 == 1999: # print every 2000 mini-batches |