CUDA не хватает памяти при тонкой настройке сегформера mit-b0 ⇐ Python

[Anonymous]

Я пытаюсь точно настроить модель сегментации mit-b0 на спутниковых изображениях для сегментации рисовых полей на RTX 2070 с 8 ГБ видеопамяти, но в начале первой эпохи у меня заканчивается память CUDA. Я считаю, что у меня возникла проблема с распределением памяти. Пожалуйста, дайте мне знать, в чем проблема. Я думаю, что смогу разместить b0 на своем графическом процессоре

Код: Выделить всё

def train_step(model: nn.Module,
dataloader: torch.utils.data.DataLoader,
optimizer: torch.optim.Optimizer,
device: torch.device):
model.train()
loss = 0.0
for i, batch in enumerate(dataloader):
pixel_values = batch["pixel_values"].to(device)
labels = batch["labels"].to(device)
optimizer.zero_grad()
z = model(pixel_values=pixel_values, labels=labels)
# logits_resized = nn.functional.interpolate(z.logits, size=labels.shape[-2:], mode="bilinear", align_corners=False)
# predicted_labels = torch.argmax(logits_resized, dim=1)
loss += z.loss
loss = loss / len(dataloader)
return loss

def eval_step(model: nn.Module,
dataloader: torch.utils.data.DataLoader,
device: torch.device):
model.eval()
loss = 0.0
with torch.inference_mode():
for i, batch in enumerate(dataloader):
pixel_values = batch["pixel_values"].to(device)
labels = batch["labels"].to(device)
z = model(pixel_values=pixel_values, labels=labels)
# logits_resized = nn.functional.interpolate(z.logits, size=labels.shape[-2:], mode="bilinear", align_corners=False)
# predicted_labels = torch.argmax(logits_resized, dim=1)
loss += z.loss
loss = loss / len(dataloader)
return loss

def train_loop(dataset_loc: str = None,
num_epochs: int = 1,
batch_size: int = 32,
num_workers: int = 1,
model_path: str = None):

train_images = os.path.join(dataset_loc, "images/train")
train_masks = os.path.join(dataset_loc, "masks/train")
list_of_train_images = os.listdir(train_images)

val_images = os.path.join(dataset_loc, "images/val")
val_masks = os.path.join(dataset_loc, "masks/val")
list_of_val_images = os.listdir(val_images)

train_transform = A.Compose([
A.HorizontalFlip(p=0.3),
A.VerticalFlip(p=0.3),
A.RandomRotate90(p=0.3),
])

image_processor = SegformerImageProcessor(do_resize=True, size={"height": 256, "width": 256})

train_dataset = RiceDataset(images=list_of_train_images,
image_folder=train_images,
mask_folder=train_masks,
transform=train_transform,
image_processor=image_processor)

val_dataset = RiceDataset(images=list_of_val_images,
image_folder=val_images,
mask_folder=val_masks,
image_processor=image_processor)

train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True)

val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False)

device = torch.device("cuda"  if torch.cuda.is_available() else "cpu")
model_checkpoint = "nvidia/mit-b0"
id2label =  {0: "outer", 1: "rice_paddy"}
label2id = {label: id for id, label in id2label.items()}
num_labels = len(id2label)

model = SegformerForSemanticSegmentation.from_pretrained(
model_checkpoint,
num_labels=num_labels,
id2label=id2label,
label2id=label2id,
ignore_mismatched_sizes=True,
reshape_last_stage=True
)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.00006)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)

best_val_loss = float("inf")
best_model_state_dict = None
for epoch in tqdm(range(num_epochs)):
train_loss = train_step(model=model,
dataloader=train_dataloader,
optimizer=optimizer,
device=device)

val_loss = eval_step(model=model,
dataloader=val_dataloader,
device=device)
print(
f"Epoch: {epoch+1} | "
f"train_loss: {train_loss:.4f} | "
f"val_loss: {val_loss:.4f} | "
)

scheduler.step()

if val_loss < best_val_loss:
best_val_loss = val_loss
best_model_state_dict = model.state_dict()

if best_model_state_dict is not None:
if model_path.endswith(".pth") or model_path.endswith(".pt"):
torch.save(best_model_state_dict, model_path)
else:
torch.save(best_model_state_dict, model_path + ".pth")
print(f"Best validation loss: {best_val_loss:.4f}")

print("DONE")

Я пытался запустить тот же код с размером пакета 2 в Google Colab, но столкнулся с той же проблемой

Подробнее здесь: https://stackoverflow.com/questions/785 ... mer-mit-b0