Spaces:
Sleeping
Sleeping
File size: 2,682 Bytes
ebae7d1 925fdc0 ebae7d1 925fdc0 ebae7d1 925fdc0 ebae7d1 925fdc0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
from torch.cuda.amp import GradScaler, autocast
# Create fake data for the main PATCDF model
batch_size = 1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Create fake data for all inputs
baseline_mri = torch.randn(batch_size, 1, 64, 64, 64, dtype=torch.float16).to(device)
clinical = torch.randn(batch_size, 2, 8, dtype=torch.float16).to(device)
datscan = torch.randn(batch_size, 4, dtype=torch.float16).to(device)
delta_t = torch.rand(batch_size, 1, dtype=torch.float16).to(device) * 5
clinical_delta = {
'updrs_motor': torch.randn(batch_size, 4, dtype=torch.float16).to(device),
'updrs_non_motor': torch.randn(batch_size, 3, dtype=torch.float16).to(device)
}
biomarker_delta = {
'datscan': torch.randn(batch_size, 4, dtype=torch.float16).to(device),
}
current_stage = torch.randint(0, 6, (batch_size,)).to(device)
targets = {
'ppse': torch.randn(batch_size, dtype=torch.float16).to(device),
'symptoms': {
'tremor': torch.randn(batch_size, dtype=torch.float16).to(device),
'rigidity': torch.randn(batch_size, dtype=torch.float16).to(device),
'bradykinesia': torch.randn(batch_size, dtype=torch.float16).to(device),
'posture': torch.randn(batch_size, dtype=torch.float16).to(device)
},
'biomarkers': {
'datscan_asymmetry': torch.randn(batch_size, dtype=torch.float16).to(device),
'dopamine_decline': torch.randn(batch_size, dtype=torch.float16).to(device)
},
'braak': torch.randint(0, 6, (batch_size,)).to(device)
}
# Instantiate the model and move it to the device, converting to half-precision
patcdf_model = PATCDF().to(device).half()
# Use a smaller batch size
batch_size = 40
# Gradient accumulation steps
accumulation_steps = 4
optimizer = torch.optim.AdamW(patcdf_model.parameters(), lr=1e-4)
for i in range(accumulation_steps):
pred_noise, pareg_outs = patcdf_model(
baseline_mri,
clinical,
datscan,
delta_t,
clinical_delta,
biomarker_delta,
current_stage,
targets=targets,
return_reg_losses=True
)
loss = F.mse_loss(pred_noise, torch.randn_like(pred_noise)) / accumulation_steps
if 'reg_loss' in pareg_outs:
loss += pareg_outs['reg_loss'] / accumulation_steps
# Scale the loss and call backward to accumulate gradients
loss.backward()
# Update the weights
optimizer.step()
optimizer.zero_grad()
# Print the output shapes
print("Predicted noise shape:", pred_noise.shape)
print("PAREG outputs:", {k: v.shape if hasattr(v, 'shape') else ({sk: sv.shape for sk, sv in v.items()} if v is not None else None) for k, v in pareg_outs.items()}) |