1. Problem Description
As the leader of an AI research team collaborating with Silkie chicken farmers, you are tasked with solving a critical challenge in traditional free-range farming. Accurate counting of livestock is crucial for both farmers and insurance companies, as factors like disease outbreaks and predator invasions can significantly impact the survival rate of these chickens in a short time. While insurance coverage helps mitigate farming risks, the claims process requires precise counting of livestock losses. Your farmers have approached your team for help in developing more accurate, automated counting systems. The challenge before your research team is to develop an optimized Silkie chicken counting model using density estimation techniques that can provide reliable counts to support both farm management and insurance processes.
Your team has access to a pretrained feature extractor for Silkie chicken images, but you'll need to design and train the density estimation decoder to create a complete counting solution. Your task is to build upon this foundation by developing an effective decoder architecture and training strategy to achieve accurate chicken counts that farmers and insurance companies can rely on.
The below figure shows an image in the dataset, as well as the corresponding true density distribution and a predicted density distribution generated by the baseline model. The total density (sum of densities across all areas) is labeled.
2. Dataset
The structure of the provided Silkie chicken image dataset is as follows:
datasets/
├── train/
│ └── A dataset with features:
│ ├── `image`: `PIL.Image` with RGB channels (3x720x1280)
│ └── `density`: a 2D array of shape 180x320
└── base.pth (Pretrained Model)
os.environ.get("DATA_PATH")/
├── test_a/
│ └── A dataset with features:
│ └── `image`: `PIL.Image` with RGB channels (3x720x1280)
└── test_b/
└── A dataset with features:
└── `image`: `PIL.Image` with RGB channels (3x720x1280)
(1) Training set location: datasets, files in this folder are used for model fine-tuning. It contains a train folder, which stores a dataset with 100 images and their corresponding density maps.
(2) Validation set (test_a) and Test set (test_b): These will be used to evaluate scores on Leaderboard A and Leaderboard B, respectively. They will be inaccessible to contestants. Only the score achieved on test set B will be used for final scoring.
(3) Datasets size:
- Training set: 100 images.
- Validation set: 100 images.
- Test set: 100 images.
(4) Validation set (test_a) and Test set (test_b) are not visible.
(5) Due to limits of computing resources, training density maps are reshaped to . NOTE the output density map can be viewed as a 2D real number matrix with shape , the sum of all matrix values is the count of chickens.
3. Task
Your task is to train your own model using the training data to predict density maps, thereby serving the purpose of chicken counting.
You may extend and optimize the given pretrained model to improve its count prediction accuracy. The pretrained model base.pth only contains the weights of the first four layers of the model (the feature extraction model), and the function load_pretrained_weights_partial in the baseline code can be used to load these partial weights into your model. You may construct a density decoder DensityDecoder and combine it with the pretrained feature extraction module to form a complete data prediction model.
class DensityDecoder(nn.Module):
def __init__(self):
#################################################
# Your code here
#################################################
def forward(self, x):
#################################################
# Your code here
#################################################
return x
You can also build your own model without the pretrained model we provided.
This task is the continuation of Satellite Weather Forecasting. A kind remind is the UNET is easily to full GPU memory without any feature engineering. Then, the GPU memory error message will be reported.
Please follow these rules to achieve a score normally:
(1) Your model must output the predicted density map.
(2) Due to limits of computing resources, your output density map should be reshaped to . This is also the shape of target density maps provided in the train dataset.
4. Submission
Please submit a submission.ipynb that includes the following components:
（1）Training Code
- Include the full training pipeline.
（2）Evaluation Code
-
Evaluate your model on the validation set and test set.
-
The output should be saved as
submission.npz. This must be a validnpzfile containing two arrayspred_aandpred_b, each with shape100x1x180x320(The evaluation script will also accept predictions in the shape of100x180x320, if you decide to squeeze the channel dimension).Any result that does not meet the specified size will be considered invalid, resulting in an assessment score of zero.
-
Each element of the density map should be no less than zero, otherwise will result in an assessment score of zero.
5. Scoring
You will be scored based on the mean relative error of your model. Relative error is defined by:
where is the true total density for the -th sample, and is the predicted total density for the -th sample.
Your final score before normalization will be calculated based on your mean relative error, as follows:
6. Baseline & Training Set
- Below you can find the baseline solution.
- The dataset is in
training_setfolder. - The highest score by the Scientific Committee for this task is 0.89, this score is used for score unification.
- The baseline score by the Scientific Committee for this task is 0.71, this score is used for score unification.
Imports
import random
import numpy as np
import torch
seed = 42
random.seed(seed) # Python built-in random
np.random.seed(seed) # NumPy
torch.manual_seed(seed) # PyTorch (CPU)
torch.cuda.manual_seed(seed) # PyTorch (single GPU)
torch.cuda.manual_seed_all(seed) # PyTorch (all GPUs)
# Ensures deterministic behavior
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from datasets import load_from_disk
import logging
from torchvision import transforms
from tqdm import tqdm
import numpy as np
import math
#Contestants should mount "counting_problem_train(v1)" datasets while creating the node.
TRAIN_PATH = "./" #Address of the training set and base.pth
# The training set is deployed automatically in the testing machine.
# You notebook can access the TRAIN_PATH even if you do not mount it along with notebook.
TRAINING_SET = TRAIN_PATH + "train" #Address of the traninig set
BASE_MODEL_PATH = TRAIN_PATH + "base.pth"#Address of the .pth file
DTYPE = torch.float32
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
scale = 100.0
Logging Utilities
def logging_level(level='info'):
str_format = '%(asctime)s - %(levelname)s: %(message)s'
if level == 'debug':
logging.basicConfig(level=logging.DEBUG, format=str_format, datefmt='%Y-%m-%d %H:%M:%S')
elif level == 'info':
logging.basicConfig(level=logging.INFO, format=str_format, datefmt='%Y-%m-%d %H:%M:%S')
return logging
class BatchLossLogger:
def __init__(self, log_interval=100):
self.losses = []
self.batch_number = 0
self.log_interval = log_interval
def log(self, loss):
self.losses.append(loss)
self.batch_number += 1
if self.batch_number % self.log_interval == 0:
logging.info(f'Batch No. {self.batch_number:7d} - loss: {loss:.6f}')
Training Your Model
Model Definition
Pretrained weights of the FeatureExtraction model is provided in base.pth in the training set, along with a function to load them.
ChickenCounting is the completed model.
class FeatureExtraction(nn.Module):
def __init__(self, in_channels=3):
super(FeatureExtraction, self).__init__()
self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=2, dilation=2)
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=2, dilation=2)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=2, dilation=2)
self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=2, dilation=2)
self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
x = self.pool4(x)
return x
def load_pretrained_weights_partial(self, weights_path, num_layers=4):
save_model = torch.load(weights_path)
partial_state_dict = {}
expected_layers = [
'feature_extraction.conv1.weight', 'feature_extraction.conv1.bias',
'feature_extraction.conv2.weight', 'feature_extraction.conv2.bias',
'feature_extraction.conv3.weight', 'feature_extraction.conv3.bias',
'feature_extraction.conv4.weight', 'feature_extraction.conv4.bias',
]
state_dict = {k.split('.', 1)[-1]: v for k, v in save_model.items() if k in expected_layers[:2 * num_layers]}
model_dict = self.state_dict()
for k in state_dict:
if k in model_dict:
partial_state_dict[k] = state_dict[k]
print(k)
self.load_state_dict(partial_state_dict)
class DensityDecoder(nn.Module): # Define your decoder model here.
def __init__(self):
super(DensityDecoder, self).__init__()
self.conv5 = nn.Conv2d(in_channels=128, out_channels=1, kernel_size=3, padding=2, dilation=2)
def forward(self, x):
x = F.relu(self.conv5(x))
return x
class ChickenCounting(nn.Module):
def __init__(self):
super(ChickenCounting, self).__init__()
self.feature_extraction = FeatureExtraction()
self.feature_decoder = DensityDecoder()
def forward(self, x):
x = self.feature_extraction(x)
x = self.feature_decoder(x)
return x
Reading the Dataset
Read the train dataset
from datasets import load_dataset
# 从 Hugging Face 加载数据集
train_dataset = load_dataset("ioaihsc/Task2_Chicken_Counting_Train2",
data_dir="train",
split="train")
image_transform = transforms.Compose([
transforms.ToTensor(),
])
def collate_fn(batch, scale=scale):
return {
"image": torch.stack([image_transform(item["image"]) for item in batch]),
"density": torch.stack([torch.tensor(item["density"], dtype=DTYPE).unsqueeze(0) * scale for item in batch])
}
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)
Run Training
#Definition of the training process
def train_chickenfcn(model, train_loader, val_loader, optimizer, scheduler, num_epochs, device, save_path):
model.train()
criterion_mse = torch.nn.MSELoss(reduction='sum').to(device)
criterion_mae = torch.nn.L1Loss(reduction='sum').to(device)
best_loss = float('inf')
print(train_loader.__len__())
for epoch in range(num_epochs):
train_loss_mse = 0.0
train_loss_mae = 0.0
train_loader_tqdm = tqdm(train_loader, desc=f'Epoch {epoch + 1}/{num_epochs}', leave=False)
for i, data in enumerate(train_loader_tqdm, 0):
inputs, targets = data["image"], data["density"]
inputs = inputs.to(device).float()
targets = targets.to(device).float()
# print(targets.shape)
# t = np.sum((targets[0] / scale).cpu().numpy().squeeze())
# print(t)
optimizer.zero_grad()
outputs = model(inputs)
loss_mse = criterion_mse(outputs, targets)
loss_mae = criterion_mae(outputs, targets)
loss_mae.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0)
optimizer.step()
train_loss_mse += loss_mse.item()
train_loss_mae += loss_mae.item()
train_loader_tqdm.set_postfix({'Train MSE Loss': loss_mse.item(), 'Train MAE Loss': loss_mae.item()})
train_loss_mse /= (len(train_loader))
train_loss_mae /= (len(train_loader))
logging.info(
f'Epoch [{epoch + 1}/{num_epochs}], Train MSE loss: {train_loss_mse:.8f}, MAE loss: {train_loss_mae:.8f}')
scheduler.step()
# Validation
model.eval()
val_loss_mse = 0.0
val_loss_mae = 0.0
val_loader_tqdm = tqdm(val_loader, desc=f'Validation Epoch {epoch + 1}/{num_epochs}', leave=False)
with torch.no_grad():
for i, data in enumerate(val_loader_tqdm, 0):
inputs, targets = data["image"], data["density"]
inputs = inputs.to(device).float()
targets = targets.to(device).float()
outputs = model(inputs)
mse_loss = criterion_mse(outputs, targets)
mae_loss = criterion_mae(outputs, targets)
val_loss_mse += mse_loss.item()
val_loss_mae += mae_loss.item()
val_loader_tqdm.set_postfix(
{'Validation MSE Loss': mse_loss.item(), 'Validation MAE Loss': mae_loss.item()})
val_loss_mse /= (len(val_loader))
val_loss_mae /= (len(val_loader))
logging.info(
f'Epoch [{epoch + 1}/{num_epochs}], Validation MSE Loss: {val_loss_mse:.8f}, MAE Loss: {val_loss_mae:.8f}')
# Save Model
if val_loss_mae < best_loss:
best_loss = val_loss_mae
torch.save(model.state_dict(), save_path)
print('Finished Training ChickenFCN')
logging = logging_level('info')
logging.debug('use debug level logging setting')
################################################################################
# Experiment Settings
################################################################################
learning_rate = 1e-4
lr_decay = 1e-5
weight_decay = 0.0001
save_path = "model.pth"
epochs = 20
# Training
model = ChickenCounting().to(DEVICE)
model.feature_extraction.load_pretrained_weights_partial(BASE_MODEL_PATH)
print('load model success')
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=1 - lr_decay)
logging.info('Begin training single view model...')
train_chickenfcn(model, train_loader, val_loader, optimizer, scheduler, epochs, DEVICE, save_path=save_path)
logging.info('Finished training single view model.')
Evaluate Model
This section validates the model on the train set, which helps contestants understand whether the model is usable and calculate the score on the train set.
# Definition of the evaluation function
def evaluate(model, val_loader, device, scale): # Function used for final scoring.
model.eval() # Set the model to evaluation mode
# Initialize metrics
mse = 0.0
mae = 0.0
predict_num = 0.0
true_num = 0.0
rate = 0.0
with torch.no_grad(): # Disable gradient calculation for inference
for i, data in enumerate(val_loader, 0):
inputs, targets = data["image"], data["density"]
inputs = inputs.to(device).float() # Move inputs to device and convert to float
targets = targets.to(device).float() # Move targets to device and convert to float
# Get the model predictions
outputs = model(inputs) / scale # Adjusting for the scaling factor
# Convert tensors to numpy for visualization and metrics calculation
inputs_np = inputs.cpu().numpy() # Convert inputs to numpy
targets_np = targets.cpu().numpy() # Convert targets to numpy
outputs_np = outputs.cpu().numpy() # Convert outputs to numpy
# imshow_res(inputs_np, targets_np, outputs_np, scale) # Uncomment to visualize results
# Calculate true and predicted sums for comparison
t = np.sum((targets[0] / scale).cpu().numpy().squeeze()) # Ground truth sum
g = np.sum(outputs.cpu().numpy().squeeze()) # Predicted sum
print(f'NO.{i} true_sum={t}, get_sum={g}, abs={abs(t - g)}, rate={abs(1 - g / t)}')
# Update metrics
predict_num += g
true_num += t
rate += abs(1 - g / t)
mae += abs(t - g)
mse += abs(t - g) * abs(t - g)
# Calculate average metrics across all batches
mae /= len(val_loader)
mse /= len(val_loader)
predict_num /= len(val_loader)
true_num /= len(val_loader)
rate /= len(val_loader)
# Log the results
logging.info(
f'test ---- Score: {math.exp(-rate):.3f}, MSE: {mse:.4f}, MAE: {mae:.4f}, Chicken_avg: {predict_num:.4f}')
return math.exp(-rate)
model.load_state_dict(torch.load(save_path, map_location=DEVICE))
model.to(DEVICE)
evaluate(model, val_loader, DEVICE, scale)
Submission
This part is to generate the result files for testing and scoring. Contestants couldn't access the validation set(test_a) and the test set(test_b) locally. Please read through the following code carefully. Make sure to following the file naming conventions.
from datasets import load_dataset
test_dataset = load_dataset("ioaihsc/Task2_Chicken_Counting_Test",
data_dir="valandtest",
split="validation")
def collate_fn(batch): # The test datasets will not provide target densities
return torch.stack([image_transform(item["image"]) for item in batch])
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn)
predictions = []
model.eval()
with torch.no_grad():
for batch in tqdm(test_loader):
outputs = model(batch.to(DEVICE)) / scale
predictions.append(outputs.cpu().numpy())
pred_a = np.concatenate(predictions, axis=0)
del test_dataset
del test_loader
del predictions
test_dataset = load_dataset("ioaihsc/Task2_Chicken_Counting_Test",
data_dir="valandtest",
split="test")
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, collate_fn=collate_fn)
predictions = []
with torch.no_grad():
for batch in tqdm(test_loader):
outputs = model(batch.to(DEVICE)) / scale
predictions.append(outputs.cpu().numpy())
pred_b = np.concatenate(predictions, axis=0)
np.savez('submission.npz', pred_a=pred_a, pred_b=pred_b) # save your submissions in `submission.npz` file with the keys `pred_a` and `pred_b`
%run ./metrics.py