<jupyter_start><jupyter_text>Load train and validation data<jupyter_code>dataset = get_dataset(dataset="fmow", download=False)
grouper = CombinatorialGrouper(dataset, ["region"])

train_data = dataset.get_subset(
    "train",
    transform=transforms.Compose(
        [transforms.ToTensor()]
    ),
)

ood_val_data = dataset.get_subset(
    "val",
    transform=transforms.Compose(
        [transforms.ToTensor()]
    ),
)

id_val_data = dataset.get_subset(
    "id_val",
    transform=transforms.Compose(
        [transforms.ToTensor()]
    ),
)

# train_dataset = SubsampledDataset(train_data, grouper)
# ood_val_dataset = SubsampledDataset(ood_val_data, grouper)
# id_val_dataset = SubsampledDataset(id_val_data, grouper)

train_dataset = torch.load("subsampled_data/train.pt")
ood_val_dataset = torch.load("subsampled_data/ood_val.pt")
id_val_dataset = torch.load("subsampled_data/id_val.pt")

train_loader = get_train_loader("standard", train_dataset, batch_size=BATCH_SIZE)
ood_val_loader   = get_train_loader("standard", ood_val_dataset, batch_size=BATCH_SIZE)
id_val_loader    = get_train_loader("standard", id_val_dataset, batch_size=BATCH_SIZE)<jupyter_output><empty_output><jupyter_text>Train and accumulate evaluation per epoch<jupyter_code>NUM_EPOCHS = 10
LEARNING_RATE = 0.0001
MODEL_PATH = "models"
NUM_CLASSES = 10
MOMENTUM = 0.9

model = Classifier(NUM_CLASSES, use_pretrained=True, freeze_encoder=True)
model.to(device)

# summary = summary(model, (3, 224, 224))

model_name = f"ERM_{NUM_CLASSES}_{NUM_EPOCHS}_SGD_{LEARNING_RATE}_{MOMENTUM}_CrossEntropy_with_unet_with_selfsup_frozen.pth"
save_name = os.path.join(MODEL_PATH, model_name)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM)

train_evolution = []
val_evolution = []
id_val_evolution = []
best_loss = sys.float_info.max

for epoch in range(NUM_EPOCHS):
	print(f"EPOCH {epoch + 1}:")
	# train
	y_true, y_pred, metadata, loss = train_step(model, train_loader, loss_fn, optimizer, device)
	train_evolution.append(build_metrics_dict(dataset, y_true, y_pred, metadata, loss))
	print(f"Train Loss: {loss}")

	# validation
	y_true, y_pred, metadata, loss = val_step(model, ood_val_loader, loss_fn, device)
	val_evolution.append(build_metrics_dict(dataset, y_true, y_pred, metadata, loss))
	print(f"OOD Loss: {loss}")

	# # save by best ood loss
	if loss < best_loss:
		best_loss = loss
		torch.save(model, save_name)

	# in distribution validation
	y_true, y_pred, metadata, loss = val_step(model, id_val_loader, loss_fn, device)
	id_val_evolution.append(build_metrics_dict(dataset, y_true, y_pred, metadata, loss))
	print(f"ID Loss: {loss}")<jupyter_output>EPOCH 1:<jupyter_text>Plot loss and accuracy per region<jupyter_code>metrics = list(train_evolution[0].keys())

for metric in metrics:
	plot_graph(metric, train_evolution, val_evolution, id_val_evolution, NUM_EPOCHS)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1164685908319185, 'fraction_numerical': 0.007130730050933786, 'mean_word_length': 6.940700808625337, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '637637', 'license': 'MIT', 'n_tokens_mistral': 1131, 'n_tokens_neox': 1096, 'n_words': 181, 'path': 'experiments/train_ERM.ipynb', 'repo_name': 'Dawlau/fMoW-classification', 'size': '2945'}