Imports
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 Dataset, DataLoader
from datasets import load_from_disk
import logging
from torchvision import transforms
from tqdm import tqdm
import numpy as np
import math
from PIL import Image
import zipfile
from typing import Optional
import random
np.random.seed(42)
random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)
TRAIN_PATH = "./"
TRAINING_SET = TRAIN_PATH + "train"
BASE_MODEL_PATH = TRAIN_PATH + "base.pth"
DTYPE = torch.float32
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
scale = 300.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}')
Traininng Your Model
Model Definition
Pretrained weights of the FeatureExtraction model is provided, along with a function to load them. And a completed UNet structure is applied.
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else:
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=False):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = (DoubleConv(n_channels, 64))
self.down1 = (Down(64, 128))
self.down2 = (Down(128, 256))
self.down3 = (Down(256, 512))
factor = 2 if bilinear else 1
self.down4 = (Down(512, 1024 // factor))
self.up1 = (Up(1024, 512 // factor, bilinear))
self.up2 = (Up(512, 256 // factor, bilinear))
# self.up2 = (Up(512, 64, bilinear))
self.up3 = (Up(256, 128 // factor, bilinear))
self.up4 = (Up(128, 64, bilinear))
self.outc = (OutConv(64, n_classes))
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = self.outc(x)
logits = F.relu(logits)
return logits
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 ChickenCounting(nn.Module):
def __init__(self):
super(ChickenCounting, self).__init__()
self.feature_extraction = FeatureExtraction()
self.unet = UNet(n_channels=128, n_classes=1)
def forward(self, x):
x = self.feature_extraction(x)
x = self.unet(x)
return x
Reading the 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=1, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)
Training
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 = 30
################################################################################
# Dataset paths
################################################################################
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)
# optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, 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
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
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