develop #4
|
@ -15,7 +15,7 @@ class UpscaleDataset(Dataset):
|
||||||
combined_df = pd.DataFrame()
|
combined_df = pd.DataFrame()
|
||||||
for parquet_file in parquet_files:
|
for parquet_file in parquet_files:
|
||||||
# Load data with chunking for memory efficiency
|
# Load data with chunking for memory efficiency
|
||||||
df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(5000)
|
df = pd.read_parquet(parquet_file, columns=['image_512', 'image_1024']).head(2500)
|
||||||
combined_df = pd.concat([combined_df, df], ignore_index=True)
|
combined_df = pd.concat([combined_df, df], ignore_index=True)
|
||||||
|
|
||||||
# Validate data format
|
# Validate data format
|
||||||
|
@ -126,7 +126,7 @@ for epoch in range(num_epochs):
|
||||||
optimizer.zero_grad()
|
optimizer.zero_grad()
|
||||||
|
|
||||||
# Use automatic mixed precision context
|
# Use automatic mixed precision context
|
||||||
with autocast(device_type=device):
|
with autocast(device_type="cuda"):
|
||||||
outputs = model(low_res)
|
outputs = model(low_res)
|
||||||
loss = criterion(outputs, high_res)
|
loss = criterion(outputs, high_res)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue